From ff0e5f5cb401bbcf6f3c9535e2fcd57474967b11 Mon Sep 17 00:00:00 2001
From: Charlie Marsh <charlie.r.marsh@gmail.com>
Date: Tue, 15 Nov 2022 12:19:22 -0500
Subject: [PATCH] Preserve scopes when checking deferred strings (#758)

---
 resources/test/fixtures/F821_5.py             | 14 +++++++++
 src/check_ast.rs                              | 31 ++++++++++++++-----
 src/linter.rs                                 |  1 +
 .../ruff__linter__tests__F821_F821_5.py.snap  | 14 +++++++++
 4 files changed, 52 insertions(+), 8 deletions(-)
 create mode 100644 resources/test/fixtures/F821_5.py
 create mode 100644 src/snapshots/ruff__linter__tests__F821_F821_5.py.snap

diff --git a/resources/test/fixtures/F821_5.py b/resources/test/fixtures/F821_5.py
new file mode 100644
index 0000000000..75bf0c3489
--- /dev/null
+++ b/resources/test/fixtures/F821_5.py
@@ -0,0 +1,14 @@
+"""Test: inner class annotation."""
+
+
+class RandomClass:
+    def random_func(self) -> "InnerClass":
+        pass
+
+
+class OuterClass:
+    class InnerClass:
+        pass
+
+    def failing_func(self) -> "InnerClass":
+        return self.InnerClass()
diff --git a/src/check_ast.rs b/src/check_ast.rs
index 27112d4ee1..d91da72137 100644
--- a/src/check_ast.rs
+++ b/src/check_ast.rs
@@ -65,7 +65,7 @@ pub struct Checker<'a> {
     scopes: Vec<Scope<'a>>,
     scope_stack: Vec<usize>,
     dead_scopes: Vec<usize>,
-    deferred_string_annotations: Vec<(Range, &'a str)>,
+    deferred_string_annotations: Vec<(Range, &'a str, Vec<usize>, Vec<usize>)>,
     deferred_annotations: Vec<(&'a Expr, Vec<usize>, Vec<usize>)>,
     deferred_functions: Vec<(&'a Stmt, Vec<usize>, Vec<usize>, VisibleScope)>,
     deferred_lambdas: Vec<(&'a Expr, Vec<usize>, Vec<usize>)>,
@@ -1042,8 +1042,12 @@ where
                 ..
             } = &expr.node
             {
-                self.deferred_string_annotations
-                    .push((Range::from_located(expr), value));
+                self.deferred_string_annotations.push((
+                    Range::from_located(expr),
+                    value,
+                    self.scope_stack.clone(),
+                    self.parent_stack.clone(),
+                ));
             } else {
                 self.deferred_annotations.push((
                     expr,
@@ -1569,8 +1573,12 @@ where
                 ..
             } => {
                 if self.in_annotation && !self.in_literal {
-                    self.deferred_string_annotations
-                        .push((Range::from_located(expr), value));
+                    self.deferred_string_annotations.push((
+                        Range::from_located(expr),
+                        value,
+                        self.scope_stack.clone(),
+                        self.parent_stack.clone(),
+                    ));
                 }
                 if self.settings.enabled.contains(&CheckCode::S104) {
                     if let Some(check) = flake8_bandit::plugins::hardcoded_bind_all_interfaces(
@@ -2128,6 +2136,7 @@ impl<'a> Checker<'a> {
             let mut import_starred = false;
             for scope_index in self.scope_stack.iter().rev() {
                 let scope = &mut self.scopes[*scope_index];
+
                 if matches!(scope.kind, ScopeKind::Class(_)) {
                     if id == "__class__" {
                         return;
@@ -2344,8 +2353,8 @@ impl<'a> Checker<'a> {
 
     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.parent_stack = parents;
             self.visit_expr(expr);
         }
     }
@@ -2354,10 +2363,14 @@ impl<'a> Checker<'a> {
     where
         'b: 'a,
     {
-        while let Some((range, expression)) = self.deferred_string_annotations.pop() {
+        let mut stacks = vec![];
+        while let Some((range, expression, scopes, parents)) =
+            self.deferred_string_annotations.pop()
+        {
             if let Ok(mut expr) = parser::parse_expression(expression, "<filename>") {
                 relocate_expr(&mut expr, range);
                 allocator.push(expr);
+                stacks.push((scopes, parents));
             } else {
                 if self.settings.enabled.contains(&CheckCode::F722) {
                     self.add_check(Check::new(
@@ -2367,7 +2380,9 @@ impl<'a> Checker<'a> {
                 }
             }
         }
-        for expr in allocator {
+        for (expr, (scopes, parents)) in allocator.iter().zip(stacks) {
+            self.scope_stack = scopes;
+            self.parent_stack = parents;
             self.visit_expr(expr);
         }
     }
diff --git a/src/linter.rs b/src/linter.rs
index cb6c3691ee..2f1098995a 100644
--- a/src/linter.rs
+++ b/src/linter.rs
@@ -456,6 +456,7 @@ mod tests {
     #[test_case(CheckCode::F821, Path::new("F821_2.py"); "F821_2")]
     #[test_case(CheckCode::F821, Path::new("F821_3.py"); "F821_3")]
     #[test_case(CheckCode::F821, Path::new("F821_4.py"); "F821_4")]
+    #[test_case(CheckCode::F821, Path::new("F821_5.py"); "F821_5")]
     #[test_case(CheckCode::F822, Path::new("F822.py"); "F822")]
     #[test_case(CheckCode::F823, Path::new("F823.py"); "F823")]
     #[test_case(CheckCode::F831, Path::new("F831.py"); "F831")]
diff --git a/src/snapshots/ruff__linter__tests__F821_F821_5.py.snap b/src/snapshots/ruff__linter__tests__F821_F821_5.py.snap
new file mode 100644
index 0000000000..f0546ccece
--- /dev/null
+++ b/src/snapshots/ruff__linter__tests__F821_F821_5.py.snap
@@ -0,0 +1,14 @@
+---
+source: src/linter.rs
+expression: checks
+---
+- kind:
+    UndefinedName: InnerClass
+  location:
+    row: 5
+    column: 29
+  end_location:
+    row: 5
+    column: 41
+  fix: ~
+