Adjust location when parsing deferred type annotations (#133)

resources/test/fixtures/F821.py

@@ -54,3 +54,5 @@ except Exception as e:
 
 
 y: int = 1
+
+x: "Bar" = 1

src/cache.rs

@@ -3,12 +3,12 @@ use std::fs::Metadata;
 use std::hash::{Hash, Hasher};
 use std::path::Path;
 
-use crate::autofix;
 use cacache::Error::EntryNotFound;
 use filetime::FileTime;
 use log::error;
 use serde::{Deserialize, Serialize};
 
+use crate::autofix;
 use crate::message::Message;
 use crate::settings::Settings;
 

src/check_ast.rs

@@ -13,6 +13,7 @@ use crate::ast_ops::{
 };
 use crate::builtins::{BUILTINS, MAGIC_GLOBALS};
 use crate::checks::{Check, CheckCode, CheckKind, Fix, RejectedCmpop};
+use crate::relocator::relocate_expr;
 use crate::settings::Settings;
 use crate::visitor::{walk_excepthandler, Visitor};
 use crate::{autofix, fixer, visitor};
@@ -34,7 +35,7 @@ struct Checker<'a> {
     scopes: Vec<Scope>,
     scope_stack: Vec<usize>,
     dead_scopes: Vec<usize>,
-    deferred_annotations: Vec<&'a str>,
+    deferred_annotations: Vec<(Location, &'a str)>,
     deferred_functions: Vec<(&'a Stmt, Vec<usize>, Vec<usize>)>,
     deferred_lambdas: Vec<(&'a Expr, Vec<usize>, Vec<usize>)>,
     // Derivative state.
@@ -759,7 +760,7 @@ where
                 value: Constant::Str(value),
                 ..
             } if self.in_annotation && !self.in_literal => {
-                self.deferred_annotations.push(value);
+                self.deferred_annotations.push((expr.location, value));
             }
             _ => {}
         };
@@ -1074,7 +1075,6 @@ impl<'a> Checker<'a> {
                     })
                     .unwrap_or_default()
             {
-                // Really need parent here.
                 self.add_binding(
                     id.to_string(),
                     Binding {
@@ -1115,8 +1115,9 @@ impl<'a> Checker<'a> {
         'b: 'a,
     {
         while !self.deferred_annotations.is_empty() {
-            let value = self.deferred_annotations.pop().unwrap();
+            let (location, expression) = self.deferred_annotations.pop().unwrap();
-            if let Ok(expr) = parser::parse_expression(value, path) {
+            if let Ok(mut expr) = parser::parse_expression(expression, path) {
+                relocate_expr(&mut expr, location);
                 allocator.push(expr);
             }
         }

src/fixer.rs

@@ -1,8 +1,8 @@
-use crate::ast_ops::SourceCodeLocator;
 use rustpython_parser::ast::{Expr, Keyword, Location};
 use rustpython_parser::lexer;
 use rustpython_parser::token::Tok;
 
+use crate::ast_ops::SourceCodeLocator;
 use crate::checks::Fix;
 
 /// Convert a location within a file (relative to `base`) to an absolute position.

src/lib.rs

@@ -13,5 +13,6 @@ pub mod linter;
 pub mod logging;
 pub mod message;
 mod pyproject;
+mod relocator;
 pub mod settings;
 mod visitor;

src/linter.rs

@@ -664,6 +664,11 @@ mod tests {
                 location: Location::new(21, 12),
                 fix: None,
             },
+            Check {
+                kind: CheckKind::UndefinedName("Bar".to_string()),
+                location: Location::new(58, 5),
+                fix: None,
+            },
         ];
         assert_eq!(actual.len(), expected.len());
         for i in 0..actual.len() {

src/relocator.rs

@@ -0,0 +1,137 @@
+use rustpython_parser::ast::{Expr, ExprKind, Keyword, Location};
+
+fn relocate_keyword(keyword: &mut Keyword, location: Location) {
+    keyword.location = location;
+    relocate_expr(&mut keyword.node.value, location);
+}
+
+/// Change an expression's location (recursively) to match a desired, fixed location.
+pub fn relocate_expr(expr: &mut Expr, location: Location) {
+    expr.location = location;
+    match &mut expr.node {
+        ExprKind::BoolOp { values, .. } => {
+            for expr in values {
+                relocate_expr(expr, location);
+            }
+        }
+        ExprKind::NamedExpr { target, value } => {
+            relocate_expr(target, location);
+            relocate_expr(value, location);
+        }
+        ExprKind::BinOp { left, right, .. } => {
+            relocate_expr(left, location);
+            relocate_expr(right, location);
+        }
+        ExprKind::UnaryOp { operand, .. } => {
+            relocate_expr(operand, location);
+        }
+        ExprKind::Lambda { body, .. } => {
+            relocate_expr(body, location);
+        }
+        ExprKind::IfExp { test, body, orelse } => {
+            relocate_expr(test, location);
+            relocate_expr(body, location);
+            relocate_expr(orelse, location);
+        }
+        ExprKind::Dict { keys, values } => {
+            for expr in keys {
+                relocate_expr(expr, location);
+            }
+            for expr in values {
+                relocate_expr(expr, location);
+            }
+        }
+        ExprKind::Set { elts } => {
+            for expr in elts {
+                relocate_expr(expr, location);
+            }
+        }
+        ExprKind::ListComp { elt, .. } => {
+            relocate_expr(elt, location);
+        }
+        ExprKind::SetComp { elt, .. } => {
+            relocate_expr(elt, location);
+        }
+        ExprKind::DictComp { key, value, .. } => {
+            relocate_expr(key, location);
+            relocate_expr(value, location);
+        }
+        ExprKind::GeneratorExp { elt, .. } => {
+            relocate_expr(elt, location);
+        }
+        ExprKind::Await { value } => relocate_expr(value, location),
+        ExprKind::Yield { value } => {
+            if let Some(expr) = value {
+                relocate_expr(expr, location);
+            }
+        }
+        ExprKind::YieldFrom { value } => relocate_expr(value, location),
+        ExprKind::Compare {
+            left, comparators, ..
+        } => {
+            relocate_expr(left, location);
+            for expr in comparators {
+                relocate_expr(expr, location);
+            }
+        }
+        ExprKind::Call {
+            func,
+            args,
+            keywords,
+        } => {
+            relocate_expr(func, location);
+            for expr in args {
+                relocate_expr(expr, location);
+            }
+            for keyword in keywords {
+                relocate_keyword(keyword, location);
+            }
+        }
+        ExprKind::FormattedValue {
+            value, format_spec, ..
+        } => {
+            relocate_expr(value, location);
+            if let Some(expr) = format_spec {
+                relocate_expr(expr, location);
+            }
+        }
+        ExprKind::JoinedStr { values } => {
+            for expr in values {
+                relocate_expr(expr, location);
+            }
+        }
+        ExprKind::Constant { .. } => {}
+        ExprKind::Attribute { value, .. } => {
+            relocate_expr(value, location);
+        }
+        ExprKind::Subscript { value, slice, .. } => {
+            relocate_expr(value, location);
+            relocate_expr(slice, location);
+        }
+        ExprKind::Starred { value, .. } => {
+            relocate_expr(value, location);
+        }
+        ExprKind::Name { .. } => {}
+        ExprKind::List { elts, .. } => {
+            for expr in elts {
+                relocate_expr(expr, location);
+            }
+        }
+        ExprKind::Tuple { elts, .. } => {
+            for expr in elts {
+                relocate_expr(expr, location);
+            }
+        }
+        ExprKind::Slice { lower, upper, step } => {
+            if let Some(expr) = lower {
+                relocate_expr(expr, location);
+            }
+            if let Some(expr) = upper {
+                relocate_expr(expr, location);
+            }
+            if let Some(expr) = step {
+                relocate_expr(expr, location);
+            }
+        }
+    }
+}