Account for typing_extensions for annotation parsing (#550)

resources/test/fixtures/F821_2.py

@@ -0,0 +1,18 @@
+"""Test: typing_extensions module imports."""
+from foo import Literal
+
+# F821 Undefined name `Model`
+x: Literal["Model"]
+
+from typing_extensions import Literal
+
+
+# OK
+x: Literal["Model"]
+
+
+import typing_extensions
+
+
+# OK
+x: typing_extensions.Literal["Model"]

src/check_ast.rs

@@ -37,13 +37,16 @@ use crate::{
 };
 
 const GLOBAL_SCOPE_INDEX: usize = 0;
-const TRACK_FROM_IMPORTS: [&str; 6] = [
+const TRACK_FROM_IMPORTS: [&str; 9] = [
     "collections",
     "collections.abc",
     "contextlib",
     "re",
     "typing",
+    "typing.io",
     "typing.re",
+    "typing_extensions",
+    "weakref",
 ];
 
 pub struct Checker<'a> {
@@ -132,6 +135,13 @@ impl<'a> Checker<'a> {
     /// Return `true` if the `Expr` is a reference to `typing.${target}`.
     pub fn match_typing_module(&self, expr: &Expr, target: &str) -> bool {
         match_name_or_attr_from_module(expr, target, "typing", self.from_imports.get("typing"))
+            || (typing::in_extensions(target)
+                && match_name_or_attr_from_module(
+                    expr,
+                    target,
+                    "typing_extensions",
+                    self.from_imports.get("typing_extensions"),
+                ))
     }
 }
 

src/linter.rs

@@ -365,6 +365,7 @@ mod tests {
     #[test_case(CheckCode::F722, Path::new("F722.py"); "F722")]
     #[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::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/python/typing.rs

@@ -3,6 +3,66 @@ use std::collections::{BTreeMap, BTreeSet};
 use once_cell::sync::Lazy;
 use rustpython_ast::{Expr, ExprKind};
 
+// See: https://pypi.org/project/typing-extensions/
+static TYPING_EXTENSIONS: Lazy<BTreeSet<&'static str>> = Lazy::new(|| {
+    BTreeSet::from([
+        "Annotated",
+        "Any",
+        "AsyncContextManager",
+        "AsyncGenerator",
+        "AsyncIterable",
+        "AsyncIterator",
+        "Awaitable",
+        "ChainMap",
+        "ClassVar",
+        "Concatenate",
+        "ContextManager",
+        "Coroutine",
+        "Counter",
+        "DefaultDict",
+        "Deque",
+        "Final",
+        "Literal",
+        "LiteralString",
+        "NamedTuple",
+        "Never",
+        "NewType",
+        "NotRequired",
+        "OrderedDict",
+        "ParamSpec",
+        "ParamSpecArgs",
+        "ParamSpecKwargs",
+        "Protocol",
+        "Required",
+        "Self",
+        "TYPE_CHECKING",
+        "Text",
+        "Type",
+        "TypeAlias",
+        "TypeGuard",
+        "TypeVar",
+        "TypeVarTuple",
+        "TypedDict",
+        "Unpack",
+        "assert_never",
+        "assert_type",
+        "clear_overloads",
+        "final",
+        "get_Type_hints",
+        "get_args",
+        "get_origin",
+        "get_overloads",
+        "is_typeddict",
+        "overload",
+        "reveal_type",
+        "runtime_checkable",
+    ])
+});
+
+pub fn in_extensions(name: &str) -> bool {
+    TYPING_EXTENSIONS.contains(name)
+}
+
 // See: https://docs.python.org/3/library/typing.html
 static IMPORTED_SUBSCRIPTS: Lazy<BTreeMap<&'static str, BTreeSet<&'static str>>> =
     Lazy::new(|| {
@@ -48,9 +108,7 @@ static IMPORTED_SUBSCRIPTS: Lazy<BTreeMap<&'static str, BTreeSet<&'static str>>>
             ("AbstractSet", "typing"),
             ("Annotated", "typing"),
             ("AsyncContextManager", "typing"),
-            ("AsyncContextManager", "typing"),
             ("AsyncGenerator", "typing"),
-            ("AsyncIterable", "typing"),
             ("AsyncIterator", "typing"),
             ("Awaitable", "typing"),
             ("BinaryIO", "typing"),
@@ -62,7 +120,6 @@ static IMPORTED_SUBSCRIPTS: Lazy<BTreeMap<&'static str, BTreeSet<&'static str>>>
             ("Concatenate", "typing"),
             ("Container", "typing"),
             ("ContextManager", "typing"),
-            ("ContextManager", "typing"),
             ("Coroutine", "typing"),
             ("Counter", "typing"),
             ("DefaultDict", "typing"),
@@ -92,7 +149,6 @@ static IMPORTED_SUBSCRIPTS: Lazy<BTreeMap<&'static str, BTreeSet<&'static str>>>
             ("TextIO", "typing"),
             ("Tuple", "typing"),
             ("Type", "typing"),
-            ("Type", "typing"),
             ("TypeGuard", "typing"),
             ("Union", "typing"),
             ("Unpack", "typing"),
@@ -104,6 +160,23 @@ static IMPORTED_SUBSCRIPTS: Lazy<BTreeMap<&'static str, BTreeSet<&'static str>>>
             // `typing.re`
             ("Match", "typing.re"),
             ("Pattern", "typing.re"),
+            // `typing_extensions`
+            ("Annotated", "typing_extensions"),
+            ("AsyncContextManager", "typing_extensions"),
+            ("AsyncGenerator", "typing_extensions"),
+            ("AsyncIterable", "typing_extensions"),
+            ("AsyncIterator", "typing_extensions"),
+            ("Awaitable", "typing_extensions"),
+            ("ChainMap", "typing_extensions"),
+            ("ClassVar", "typing_extensions"),
+            ("Concatenate", "typing_extensions"),
+            ("ContextManager", "typing_extensions"),
+            ("Coroutine", "typing_extensions"),
+            ("Counter", "typing_extensions"),
+            ("DefaultDict", "typing_extensions"),
+            ("Deque", "typing_extensions"),
+            ("Literal", "typing_extensions"),
+            ("Type", "typing_extensions"),
             // `weakref`
             ("WeakKeyDictionary", "weakref"),
             ("WeakSet", "weakref"),
@@ -187,7 +260,8 @@ pub fn match_annotated_subscript(
 }
 
 /// Returns `true` if `Expr` represents a reference to a typing object with a
-/// PEP 585 built-in.
+/// PEP 585 built-in. Note that none of the PEP 585 built-ins are in
+/// `typing_extensions`.
 pub fn is_pep585_builtin(expr: &Expr, typing_imports: Option<&BTreeSet<&str>>) -> bool {
     match &expr.node {
         ExprKind::Attribute { attr, value, .. } => {

src/snapshots/ruff__linter__tests__F821_F821_2.py.snap

@@ -0,0 +1,14 @@
+---
+source: src/linter.rs
+expression: checks
+---
+- kind:
+    UndefinedName: Model
+  location:
+    row: 5
+    column: 11
+  end_location:
+    row: 5
+    column: 18
+  fix: ~
+