Fix forward annotations (#41)

README.md

@@ -125,10 +125,13 @@ ruff is distributed on [PyPI](https://pypi.org/project/ruff/), and published via
 For now, releases are cut and published manually:
 
 ```shell
-for PYTHON_VERSION in 3.7 3.8 3.9 3.10
+for TARGET in x86_64-apple-darwin aarch64-apple-darwin
 do
-  maturin publish --username crmarsh --skip-existing --target x86_64-apple-darwin -i /usr/local/opt/python@${PYTHON_VERSION}/libexec/bin/python
+  maturin publish --username crmarsh --skip-existing --target ${TARGET} -i \
-  maturin publish --username crmarsh --skip-existing --target aarch64-apple-darwin -i /usr/local/opt/python@${PYTHON_VERSION}/libexec/bin/python
+    /usr/local/opt/python@3.7/libexec/bin/python \
+    /usr/local/opt/python@3.8/libexec/bin/python \
+    /usr/local/opt/python@3.9/libexec/bin/python \
+    /usr/local/opt/python@3.10/libexec/bin/python
 done
 ```
 

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,7 +46,9 @@ struct Checker<'a> {
     checks: Vec<Check>,
     scopes: Vec<Scope>,
     dead_scopes: Vec<Scope>,
+    deferred: Vec<String>,
     in_f_string: bool,
+    in_annotation: bool,
 }
 
 impl Checker<'_> {
@@ -55,7 +58,9 @@ impl Checker<'_> {
             checks: vec![],
             scopes: vec![],
             dead_scopes: vec![],
+            deferred: vec![],
             in_f_string: false,
+            in_annotation: false,
         }
     }
 }
@@ -252,6 +257,13 @@ impl Visitor for Checker<'_> {
         }
     }
 
+    fn visit_annotation(&mut self, expr: &Expr) {
+        let initial = self.in_annotation;
+        self.in_annotation = true;
+        self.visit_expr(expr);
+        self.in_annotation = initial;
+    }
+
     fn visit_expr(&mut self, expr: &Expr) {
         let initial = self.in_f_string;
         match &expr.node {
@@ -285,6 +297,10 @@ impl Visitor for Checker<'_> {
                 }
                 self.in_f_string = true;
             }
+            ExprKind::Constant {
+                value: Constant::Str(value),
+                ..
+            } if self.in_annotation => self.deferred.push(value.to_string()),
             _ => {}
         };
 
@@ -365,6 +381,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 +425,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 +452,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)
     }
 }