Add support for from __future__ import annotations (#189)

resources/test/fixtures/future_annotations.py

@@ -0,0 +1,27 @@
+from __future__ import annotations
+
+from dataclasses import dataclass
+
+from models import Fruit, Nut
+
+
+@dataclass
+class Foo:
+    x: int
+    y: int
+
+    @classmethod
+    def a(cls) -> Foo:
+        return cls(x=0, y=0)
+
+    @classmethod
+    def b(cls) -> "Foo":
+        return cls(x=0, y=0)
+
+    @classmethod
+    def c(cls) -> Bar:
+        return cls(x=0, y=0)
+
+    @classmethod
+    def d(cls) -> Fruit:
+        return cls(x=0, y=0)

src/check_ast.rs

@@ -36,7 +36,8 @@ struct Checker<'a> {
     scopes: Vec<Scope>,
     scope_stack: Vec<usize>,
     dead_scopes: Vec<usize>,
-    deferred_annotations: Vec<(Location, &'a str)>,
+    deferred_string_annotations: Vec<(Location, &'a str)>,
+    deferred_annotations: Vec<(&'a Expr, Vec<usize>, Vec<usize>)>,
     deferred_functions: Vec<(&'a Stmt, Vec<usize>, Vec<usize>)>,
     deferred_lambdas: Vec<(&'a Expr, Vec<usize>, Vec<usize>)>,
     deferred_assignments: Vec<usize>,
@@ -47,6 +48,7 @@ struct Checker<'a> {
     seen_non_import: bool,
     seen_docstring: bool,
     futures_allowed: bool,
+    annotations_future_enabled: bool,
 }
 
 impl<'a> Checker<'a> {
@@ -67,6 +69,7 @@ impl<'a> Checker<'a> {
             scopes: vec![],
             scope_stack: vec![],
             dead_scopes: vec![],
+            deferred_string_annotations: vec![],
             deferred_annotations: vec![],
             deferred_functions: vec![],
             deferred_lambdas: vec![],
@@ -77,6 +80,7 @@ impl<'a> Checker<'a> {
             seen_non_import: false,
             seen_docstring: false,
             futures_allowed: true,
+            annotations_future_enabled: false,
         }
     }
 }
@@ -423,6 +427,10 @@ where
                             },
                         );
 
+                        if alias.node.name == "annotations" {
+                            self.annotations_future_enabled = true;
+                        }
+
                         if self.settings.select.contains(&CheckCode::F407)
                             && !ALL_FEATURE_NAMES.contains(&alias.node.name.deref())
                         {
@@ -581,6 +589,17 @@ where
         let prev_in_literal = self.in_literal;
         let prev_in_annotation = self.in_annotation;
 
+        // Important:
+        if self.in_annotation && self.annotations_future_enabled {
+            self.deferred_annotations.push((
+                expr,
+                self.scope_stack.clone(),
+                self.parent_stack.clone(),
+            ));
+            visitor::walk_expr(self, expr);
+            return;
+        }
+
         // Pre-visit.
         match &expr.node {
             ExprKind::Subscript { value, .. } => {
@@ -723,7 +742,8 @@ where
                 value: Constant::Str(value),
                 ..
             } if self.in_annotation && !self.in_literal => {
-                self.deferred_annotations.push((expr.location, value));
+                self.deferred_string_annotations
+                    .push((expr.location, value));
             }
             ExprKind::GeneratorExp { .. }
             | ExprKind::ListComp { .. }
@@ -1195,11 +1215,19 @@ impl<'a> Checker<'a> {
         }
     }
 
-    fn check_deferred_annotations<'b>(&mut self, path: &str, allocator: &'b mut Vec<Expr>)
+    fn check_deferred_annotations(&mut self) {
+        while let Some((expr, scopes, parents)) = self.deferred_annotations.pop() {
+            self.parent_stack = parents;
+            self.scope_stack = scopes;
+            self.visit_expr(expr);
+        }
+    }
+
+    fn check_deferred_string_annotations<'b>(&mut self, path: &str, allocator: &'b mut Vec<Expr>)
     where
         'b: 'a,
     {
-        while let Some((location, expression)) = self.deferred_annotations.pop() {
+        while let Some((location, expression)) = self.deferred_string_annotations.pop() {
             if let Ok(mut expr) = parser::parse_expression(expression, path) {
                 relocate_expr(&mut expr, location);
                 allocator.push(expr);
@@ -1344,8 +1372,9 @@ pub fn check_ast(
     checker.check_deferred_functions();
     checker.check_deferred_lambdas();
     checker.check_deferred_assignments();
+    checker.check_deferred_annotations();
     let mut allocator = vec![];
-    checker.check_deferred_annotations(path, &mut allocator);
+    checker.check_deferred_string_annotations(path, &mut allocator);
 
     // Reset the scope to module-level, and check all consumed scopes.
     checker.scope_stack = vec![GLOBAL_SCOPE_INDEX];

src/linter.rs

@@ -1599,4 +1599,36 @@ mod tests {
 
         Ok(())
     }
+
+    #[test]
+    fn future_annotations() -> Result<()> {
+        let mut actual = check_path(
+            Path::new("./resources/test/fixtures/future_annotations.py"),
+            &settings::Settings {
+                line_length: 88,
+                exclude: vec![],
+                select: BTreeSet::from([CheckCode::F401, CheckCode::F821]),
+            },
+            &fixer::Mode::Generate,
+        )?;
+        actual.sort_by_key(|check| check.location);
+        let expected = vec![
+            Check {
+                kind: CheckKind::UnusedImport("models.Nut".to_string()),
+                location: Location::new(5, 1),
+                fix: None,
+            },
+            Check {
+                kind: CheckKind::UndefinedName("Bar".to_string()),
+                location: Location::new(22, 19),
+                fix: None,
+            },
+        ];
+        assert_eq!(actual.len(), expected.len());
+        for i in 0..actual.len() {
+            assert_eq!(actual[i], expected[i]);
+        }
+
+        Ok(())
+    }
 }