Fix forward annotations

resources/test/src/F401.py

@@ -3,11 +3,12 @@ import functools
 from collections import (
     Counter,
     OrderedDict,
+    namedtuple,
 )
 
 
 class X:
-    def a(self):
+    def a(self) -> "namedtuple":
         x = os.environ["1"]
         y = Counter()
         return X

src/check_ast.rs

@@ -1,8 +1,9 @@
 use std::collections::{BTreeMap, BTreeSet};
 
 use rustpython_parser::ast::{
-    Arg, Arguments, Expr, ExprContext, ExprKind, Location, Stmt, StmtKind, Suite,
+    Arg, Arguments, Constant, Expr, ExprContext, ExprKind, Location, Stmt, StmtKind, Suite,
 };
+use rustpython_parser::parser;
 
 use crate::check_ast::ScopeKind::{Class, Function, Generator, Module};
 use crate::checks::{Check, CheckCode, CheckKind};
@@ -45,6 +46,7 @@ struct Checker<'a> {
     checks: Vec<Check>,
     scopes: Vec<Scope>,
     dead_scopes: Vec<Scope>,
+    deferred: Vec<String>,
     in_f_string: bool,
 }
 
@@ -55,6 +57,7 @@ impl Checker<'_> {
             checks: vec![],
             scopes: vec![],
             dead_scopes: vec![],
+            deferred: vec![],
             in_f_string: false,
         }
     }
@@ -252,6 +255,18 @@ impl Visitor for Checker<'_> {
         }
     }
 
+    fn visit_annotation(&mut self, expr: &Expr) {
+        if let ExprKind::Constant {
+            value: Constant::Str(value),
+            ..
+        } = &expr.node
+        {
+            self.deferred.push(value.to_string());
+        } else {
+            visitor::walk_expr(self, expr);
+        }
+    }
+
     fn visit_expr(&mut self, expr: &Expr) {
         let initial = self.in_f_string;
         match &expr.node {
@@ -365,6 +380,7 @@ impl Checker<'_> {
     fn add_binding(&mut self, binding: Binding) {
         // TODO(charlie): Don't treat annotations as assignments if there is an existing value.
         let scope = self.scopes.last_mut().expect("No current scope found.");
+
         scope.values.insert(
             binding.name.clone(),
             match scope.values.get(&binding.name) {
@@ -408,6 +424,14 @@ impl Checker<'_> {
         }
     }
 
+    fn check_deferred(&mut self, path: &str) {
+        for value in self.deferred.clone() {
+            if let Ok(expr) = &parser::parse_expression(&value, path) {
+                self.visit_expr(expr);
+            }
+        }
+    }
+
     fn check_dead_scopes(&mut self) {
         if self.settings.select.contains(&CheckCode::F401) {
             // TODO(charlie): Handle `__all__`.
@@ -427,15 +451,18 @@ impl Checker<'_> {
     }
 }
 
-pub fn check_ast(python_ast: &Suite, settings: &Settings) -> Vec<Check> {
+pub fn check_ast(python_ast: &Suite, settings: &Settings, path: &str) -> Vec<Check> {
     let mut checker = Checker::new(settings);
     checker.push_scope(Scope {
         kind: Module,
         values: BTreeMap::new(),
     });
+
     for stmt in python_ast {
         checker.visit_stmt(stmt);
     }
+    checker.check_deferred(path);
+
     checker.pop_scope();
     checker.check_dead_scopes();
     checker.checks

src/linter.rs

@@ -30,8 +30,9 @@ pub fn check_path(path: &Path, settings: &Settings, mode: &cache::Mode) -> Resul
         .iter()
         .any(|check_code| matches!(check_code.lint_source(), LintSource::AST))
     {
-        let python_ast = parser::parse_program(&contents, &path.to_string_lossy())?;
+        let path = path.to_string_lossy();
-        checks.extend(check_ast(&python_ast, settings));
+        let python_ast = parser::parse_program(&contents, &path)?;
+        checks.extend(check_ast(&python_ast, settings, &path));
     }
 
     // Run the lines-based checks.

src/visitor.rs

@@ -8,6 +8,9 @@ pub trait Visitor {
     fn visit_stmt(&mut self, stmt: &Stmt) {
         walk_stmt(self, stmt);
     }
+    fn visit_annotation(&mut self, expr: &Expr) {
+        walk_expr(self, expr);
+    }
     fn visit_expr(&mut self, expr: &Expr) {
         walk_expr(self, expr);
     }
@@ -80,7 +83,7 @@ pub fn walk_stmt<V: Visitor + ?Sized>(visitor: &mut V, stmt: &Stmt) {
                 visitor.visit_expr(expr)
             }
             for expr in returns {
-                visitor.visit_expr(expr);
+                visitor.visit_annotation(expr);
             }
         }
         StmtKind::AsyncFunctionDef {
@@ -100,7 +103,7 @@ pub fn walk_stmt<V: Visitor + ?Sized>(visitor: &mut V, stmt: &Stmt) {
                 visitor.visit_expr(expr)
             }
             for expr in returns {
-                visitor.visit_expr(expr);
+                visitor.visit_annotation(expr);
             }
         }
         StmtKind::ClassDef {
@@ -152,7 +155,7 @@ pub fn walk_stmt<V: Visitor + ?Sized>(visitor: &mut V, stmt: &Stmt) {
             ..
         } => {
             visitor.visit_expr(target);
-            visitor.visit_expr(annotation);
+            visitor.visit_annotation(annotation);
             if let Some(expr) = value {
                 visitor.visit_expr(expr)
             }
@@ -513,7 +516,7 @@ pub fn walk_arguments<V: Visitor + ?Sized>(visitor: &mut V, arguments: &Argument
 
 pub fn walk_arg<V: Visitor + ?Sized>(visitor: &mut V, arg: &Arg) {
     if let Some(expr) = &arg.node.annotation {
-        visitor.visit_expr(expr)
+        visitor.visit_annotation(expr)
     }
 }