Confine subscript annotation checks to `ExprContext::Load` (#583)

resources/test/fixtures/F821_3.py

@@ -0,0 +1,14 @@
+"""Test (edge case): accesses of object named `dict`."""
+
+# OK
+dict = {"parameters": {}}
+dict["parameters"]["userId"] = "userId"
+dict["parameters"]["itemId"] = "itemId"
+del dict["parameters"]["itemId"]
+
+# F821 Undefined name `key`
+# F821 Undefined name `value`
+x: dict["key", "value"]
+
+# OK
+x: dict[str, str]

src/check_ast.rs

@@ -79,6 +79,7 @@ pub struct Checker<'a> {
     in_f_string: Option<Range>,
     in_annotation: bool,
     in_literal: bool,
+    in_subscript: bool,
     seen_import_boundary: bool,
     futures_allowed: bool,
     annotations_future_enabled: bool,
@@ -118,6 +119,7 @@ impl<'a> Checker<'a> {
             in_f_string: Default::default(),
             in_annotation: Default::default(),
             in_literal: Default::default(),
+            in_subscript: Default::default(),
             seen_import_boundary: Default::default(),
             futures_allowed: true,
             annotations_future_enabled: Default::default(),
@@ -1488,36 +1490,56 @@ where
                 }
             }
             ExprKind::Subscript { value, slice, ctx } => {
-                match typing::match_annotated_subscript(value, &self.from_imports) {
-                    Some(subscript) => match subscript {
-                        // Ex) Optional[int]
-                        SubscriptKind::AnnotatedSubscript => {
-                            self.visit_expr(value);
-                            self.visit_annotation(slice);
-                            self.visit_expr_context(ctx);
-                        }
-                        // Ex) Annotated[int, "Hello, world!"]
-                        SubscriptKind::PEP593AnnotatedSubscript => {
-                            // First argument is a type (including forward references); the rest are
-                            // arbitrary Python objects.
-                            self.visit_expr(value);
-                            if let ExprKind::Tuple { elts, ctx } = &slice.node {
-                                if let Some(expr) = elts.first() {
-                                    self.visit_expr(expr);
-                                    self.in_annotation = false;
-                                    for expr in elts.iter().skip(1) {
-                                        self.visit_expr(expr);
-                                    }
-                                    self.in_annotation = true;
+                // Only allow annotations in `ExprContext::Load`. If we have, e.g.,
+                // `obj["foo"]["bar"]`, we need to avoid treating the `obj["foo"]`
+                // portion as an annotation, despite having `ExprContext::Load`. Thus, we track
+                // the `ExprContext` at the top-level.
+                let prev_in_subscript = self.in_subscript;
+                if self.in_subscript {
+                    visitor::walk_expr(self, expr);
+                } else if matches!(ctx, ExprContext::Store | ExprContext::Del) {
+                    self.in_subscript = true;
+                    visitor::walk_expr(self, expr);
+                } else {
+                    self.in_subscript = true;
+                    match typing::match_annotated_subscript(value, &self.from_imports) {
+                        Some(subscript) => {
+                            match subscript {
+                                // Ex) Optional[int]
+                                SubscriptKind::AnnotatedSubscript => {
+                                    self.visit_expr(value);
+                                    self.visit_annotation(slice);
                                     self.visit_expr_context(ctx);
                                 }
-                            } else {
-                                error!("Found non-ExprKind::Tuple argument to PEP 593 Annotation.")
+                                // Ex) Annotated[int, "Hello, world!"]
+                                SubscriptKind::PEP593AnnotatedSubscript => {
+                                    // First argument is a type (including forward references); the
+                                    // rest are arbitrary Python
+                                    // objects.
+                                    self.visit_expr(value);
+                                    if let ExprKind::Tuple { elts, ctx } = &slice.node {
+                                        if let Some(expr) = elts.first() {
+                                            self.visit_expr(expr);
+                                            self.in_annotation = false;
+                                            for expr in elts.iter().skip(1) {
+                                                self.visit_expr(expr);
+                                            }
+                                            self.in_annotation = true;
+                                            self.visit_expr_context(ctx);
+                                        }
+                                    } else {
+                                        error!(
+                                            "Found non-ExprKind::Tuple argument to PEP 593 \
+                                             Annotation."
+                                        )
+                                    }
+                                }
                             }
                         }
-                    },
-                    None => visitor::walk_expr(self, expr),
+                        None => visitor::walk_expr(self, expr),
+                    }
                 }
+                self.in_subscript = prev_in_subscript;
             }
             _ => visitor::walk_expr(self, expr),
         }

src/linter.rs

@@ -366,6 +366,7 @@ mod tests {
     #[test_case(CheckCode::F821, Path::new("F821_0.py"); "F821_0")]
     #[test_case(CheckCode::F821, Path::new("F821_1.py"); "F821_1")]
     #[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::F822, Path::new("F822.py"); "F822")]
     #[test_case(CheckCode::F823, Path::new("F823.py"); "F823")]
     #[test_case(CheckCode::F831, Path::new("F831.py"); "F831")]

src/snapshots/ruff__linter__tests__F821_F821_3.py.snap

@@ -0,0 +1,23 @@
+---
+source: src/linter.rs
+expression: checks
+---
+- kind:
+    UndefinedName: key
+  location:
+    row: 11
+    column: 8
+  end_location:
+    row: 11
+    column: 13
+  fix: ~
+- kind:
+    UndefinedName: value
+  location:
+    row: 11
+    column: 15
+  end_location:
+    row: 11
+    column: 22
+  fix: ~
+