From df438ba0513d0688e85a159fc7a84e5f5c32c327 Mon Sep 17 00:00:00 2001
From: Adrian Garcia Badaracco <1755071+adriangb@users.noreply.github.com>
Date: Thu, 6 Oct 2022 08:16:07 -0500
Subject: [PATCH] Support PEP 593 annotations (#333)

---
 resources/test/fixtures/F821.py              | 33 +++++++++
 src/check_ast.rs                             | 72 ++++++++++++++++----
 src/python/typing.rs                         |  5 ++
 src/snapshots/ruff__linter__tests__f821.snap | 27 ++++++++
 4 files changed, 124 insertions(+), 13 deletions(-)

diff --git a/resources/test/fixtures/F821.py b/resources/test/fixtures/F821.py
index 7d4d60acc811e..7adf0684abb12 100644
--- a/resources/test/fixtures/F821.py
+++ b/resources/test/fixtures/F821.py
@@ -89,3 +89,36 @@ def update_tomato():
     f'B'
     f'{B}'
 )
+
+
+from typing import Annotated, Literal
+
+
+def arbitrary_callable() -> None:
+    ...
+
+
+class PEP593Test:
+    field: Annotated[
+        int,
+        "base64",
+        arbitrary_callable(),
+        123,
+        (1, 2, 3),
+    ]
+    field_with_stringified_type: Annotated[
+        "PEP593Test",
+        123,
+    ]
+    field_with_undefined_stringified_type: Annotated[
+        "PEP593Test123",
+        123,
+    ]
+    field_with_nested_subscript: Annotated[
+        dict[Literal["foo"], str],
+        123,
+    ]
+    field_with_undefined_nested_subscript: Annotated[
+        dict["foo", "bar"],  # Expected to fail as undefined.
+        123,
+    ]
diff --git a/src/check_ast.rs b/src/check_ast.rs
index e2f214ad83bf8..28950245c9163 100644
--- a/src/check_ast.rs
+++ b/src/check_ast.rs
@@ -108,11 +108,32 @@ fn match_name_or_attr(expr: &Expr, target: &str) -> bool {
     }
 }
 
-fn is_annotated_subscript(expr: &Expr) -> bool {
+enum SubscriptKind {
+    AnnotatedSubscript,
+    PEP593AnnotatedSubscript,
+}
+
+fn match_annotated_subscript(expr: &Expr) -> Option<SubscriptKind> {
     match &expr.node {
-        ExprKind::Attribute { attr, .. } => typing::is_annotated_subscript(attr),
-        ExprKind::Name { id, .. } => typing::is_annotated_subscript(id),
-        _ => false,
+        ExprKind::Attribute { attr, .. } => {
+            if typing::is_annotated_subscript(attr) {
+                Some(SubscriptKind::AnnotatedSubscript)
+            } else if typing::is_pep593_annotated_subscript(attr) {
+                Some(SubscriptKind::PEP593AnnotatedSubscript)
+            } else {
+                None
+            }
+        }
+        ExprKind::Name { id, .. } => {
+            if typing::is_annotated_subscript(id) {
+                Some(SubscriptKind::AnnotatedSubscript)
+            } else if typing::is_pep593_annotated_subscript(id) {
+                Some(SubscriptKind::PEP593AnnotatedSubscript)
+            } else {
+                None
+            }
+        }
+        _ => None,
     }
 }
 
@@ -862,9 +883,11 @@ where
             ExprKind::Constant {
                 value: Constant::Str(value),
                 ..
-            } if self.in_annotation && !self.in_literal => {
-                self.deferred_string_annotations
-                    .push((Range::from_located(expr), value));
+            } => {
+                if self.in_annotation && !self.in_literal {
+                    self.deferred_string_annotations
+                        .push((Range::from_located(expr), value));
+                }
             }
             ExprKind::Lambda { args, .. } => {
                 // Visit the arguments, but avoid the body, which will be deferred.
@@ -1015,12 +1038,35 @@ where
                 }
             }
             ExprKind::Subscript { value, slice, ctx } => {
-                if is_annotated_subscript(value) {
-                    self.visit_expr(value);
-                    self.visit_annotation(slice);
-                    self.visit_expr_context(ctx);
-                } else {
-                    visitor::walk_expr(self, expr);
+                match match_annotated_subscript(value) {
+                    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;
+                                    self.visit_expr_context(ctx);
+                                }
+                            } else {
+                                error!("Found non-ExprKind::Tuple argument to PEP 593 Annotation.")
+                            }
+                        }
+                    },
+                    None => visitor::walk_expr(self, expr),
                 }
             }
             _ => visitor::walk_expr(self, expr),
diff --git a/src/python/typing.rs b/src/python/typing.rs
index 7ea9d464ffee5..171d5509e1513 100644
--- a/src/python/typing.rs
+++ b/src/python/typing.rs
@@ -7,6 +7,7 @@ static ANNOTATED_SUBSCRIPTS: Lazy<BTreeSet<&'static str>> = Lazy::new(|| {
         "AbstractAsyncContextManager",
         "AbstractContextManager",
         "AbstractSet",
+        // "Annotated",
         "AsyncContextManager",
         "AsyncGenerator",
         "AsyncIterable",
@@ -87,3 +88,7 @@ static ANNOTATED_SUBSCRIPTS: Lazy<BTreeSet<&'static str>> = Lazy::new(|| {
 pub fn is_annotated_subscript(name: &str) -> bool {
     ANNOTATED_SUBSCRIPTS.contains(name)
 }
+
+pub fn is_pep593_annotated_subscript(name: &str) -> bool {
+    name == "Annotated"
+}
diff --git a/src/snapshots/ruff__linter__tests__f821.snap b/src/snapshots/ruff__linter__tests__f821.snap
index cc94c3fa3440a..fa5c0624b4fd4 100644
--- a/src/snapshots/ruff__linter__tests__f821.snap
+++ b/src/snapshots/ruff__linter__tests__f821.snap
@@ -74,4 +74,31 @@ expression: checks
     row: 89
     column: 9
   fix: ~
+- kind:
+    UndefinedName: PEP593Test123
+  location:
+    row: 114
+    column: 10
+  end_location:
+    row: 114
+    column: 24
+  fix: ~
+- kind:
+    UndefinedName: foo
+  location:
+    row: 122
+    column: 15
+  end_location:
+    row: 122
+    column: 19
+  fix: ~
+- kind:
+    UndefinedName: bar
+  location:
+    row: 122
+    column: 22
+  end_location:
+    row: 122
+    column: 26
+  fix: ~