diff --git a/resources/test/fixtures/future_annotations.py b/resources/test/fixtures/future_annotations.py new file mode 100644 index 0000000000..8d43d5bc35 --- /dev/null +++ b/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) diff --git a/src/check_ast.rs b/src/check_ast.rs index 9de7c807fd..da06f01894 100644 --- a/src/check_ast.rs +++ b/src/check_ast.rs @@ -36,7 +36,8 @@ struct Checker<'a> { scopes: Vec, scope_stack: Vec, dead_scopes: Vec, - deferred_annotations: Vec<(Location, &'a str)>, + deferred_string_annotations: Vec<(Location, &'a str)>, + deferred_annotations: Vec<(&'a Expr, Vec, Vec)>, deferred_functions: Vec<(&'a Stmt, Vec, Vec)>, deferred_lambdas: Vec<(&'a Expr, Vec, Vec)>, deferred_assignments: Vec, @@ -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) + 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) 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]; diff --git a/src/linter.rs b/src/linter.rs index ce15561e76..3b2f07112e 100644 --- a/src/linter.rs +++ b/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(()) + } }