check_new_syntax.py: check for any import of deprecated contextlib aliases (#6371)

tests/check_new_syntax.py

@@ -4,19 +4,14 @@ import ast
 import sys
 from itertools import chain
 from pathlib import Path
-from typing import TypedDict
-
-
-class ModuleImportsDict(TypedDict):
-    set_from_collections_abc: bool
-    context_manager_from_typing: bool
-    async_context_manager_from_typing: bool
-
 
 STUBS_SUPPORTING_PYTHON_2 = frozenset(
-    {path.parent for path in Path("stubs").rglob("METADATA.toml") if "python2 = true" in path.read_text().splitlines()}
+    path.parent for path in Path("stubs").rglob("METADATA.toml") if "python2 = true" in path.read_text().splitlines()
 )
 
+CONTEXT_MANAGER_ALIASES = {"ContextManager": "AbstractContextManager", "AsyncContextManager": "AbstractAsyncContextManager"}
+CONTEXTLIB_ALIAS_ALLOWLIST = frozenset({Path("stdlib/contextlib.pyi"), Path("stdlib/typing_extensions.pyi")})
+
 
 def check_new_syntax(tree: ast.AST, path: Path) -> list[str]:
     errors = []
@@ -30,9 +25,12 @@ def check_new_syntax(tree: ast.AST, path: Path) -> list[str]:
     def is_dotdotdot(node: ast.AST) -> bool:
         return isinstance(node, ast.Constant) and node.s is Ellipsis
 
+    def add_contextlib_alias_error(node: ast.ImportFrom | ast.Attribute, alias: str) -> None:
+        errors.append(f"{path}:{node.lineno}: Use `contextlib.{CONTEXT_MANAGER_ALIASES[alias]}` instead of `typing.{alias}`")
+
     class OldSyntaxFinder(ast.NodeVisitor):
-        def __init__(self, module_imports_info_dict: ModuleImportsDict) -> None:
-            self.module_imports_info_dict = module_imports_info_dict
+        def __init__(self, *, set_from_collections_abc: bool) -> None:
+            self.set_from_collections_abc = set_from_collections_abc
 
         def visit_Subscript(self, node: ast.Subscript) -> None:
             if isinstance(node.value, ast.Name):
@@ -45,7 +43,7 @@ def check_new_syntax(tree: ast.AST, path: Path) -> list[str]:
                 if node.value.id in {"List", "FrozenSet"}:
                     new_syntax = f"{node.value.id.lower()}[{ast.unparse(node.slice)}]"
                     errors.append(f"{path}:{node.lineno}: Use built-in generics, e.g. `{new_syntax}`")
-                if not self.module_imports_info_dict["set_from_collections_abc"] and node.value.id == "Set":
+                if not self.set_from_collections_abc and node.value.id == "Set":
                     new_syntax = f"set[{ast.unparse(node.slice)}]"
                     errors.append(f"{path}:{node.lineno}: Use built-in generics, e.g. `{new_syntax}`")
                 if node.value.id == "Deque":
@@ -66,23 +64,6 @@ def check_new_syntax(tree: ast.AST, path: Path) -> list[str]:
                 ):
                     new_syntax = f"tuple[{unparse_without_tuple_parens(node.slice)}]"
                     errors.append(f"{path}:{node.lineno}: Use built-in generics, e.g. `{new_syntax}`")
-                if not python_2_support_required:
-                    if self.module_imports_info_dict["context_manager_from_typing"] and node.value.id == "ContextManager":
-                        new_syntax = f"contextlib.AbstractContextManager[{ast.unparse(node.slice)}]"
-                        errors.append(
-                            f"{path}:{node.lineno}: Use `contextlib.AbstractContextManager` instead of `typing.ContextManager`, "
-                            f"e.g. `{new_syntax}`"
-                        )
-                    if (
-                        self.module_imports_info_dict["async_context_manager_from_typing"]
-                        and node.value.id == "AsyncContextManager"
-                    ):
-                        new_syntax = f"contextlib.AbstractAsyncContextManager[{ast.unparse(node.slice)}]"
-                        errors.append(
-                            f"{path}:{node.lineno}: "
-                            f"Use `contextlib.AbstractAsyncContextManager` instead of `typing.AsyncContextManager`, "
-                            f"e.g. `{new_syntax}`"
-                        )
 
             self.generic_visit(node)
 
@@ -92,32 +73,33 @@ def check_new_syntax(tree: ast.AST, path: Path) -> list[str]:
     # TODO: can use built-in generics in type aliases
     class AnnotationFinder(ast.NodeVisitor):
         def __init__(self) -> None:
-            self.module_imports_info_dict = ModuleImportsDict(
-                set_from_collections_abc=False, context_manager_from_typing=False, async_context_manager_from_typing=False
-            )
+            self.set_from_collections_abc = False
 
         def old_syntax_finder(self) -> OldSyntaxFinder:
-            """Convenience method to create an `OldSyntaxFinder` method with the correct state"""
-            return OldSyntaxFinder(self.module_imports_info_dict)
+            """Convenience method to create an `OldSyntaxFinder` instance with the correct state"""
+            return OldSyntaxFinder(set_from_collections_abc=self.set_from_collections_abc)
 
         def visit_ImportFrom(self, node: ast.ImportFrom) -> None:
             if node.module == "collections.abc":
-                classes_from_collections_abc = node.names
-                if any(cls.name == "Set" for cls in classes_from_collections_abc):
-                    self.module_imports_info_dict["set_from_collections_abc"] = True
+                imported_classes = node.names
+                if any(cls.name == "Set" for cls in imported_classes):
+                    self.set_from_collections_abc = True
 
-            elif node.module == "typing":
-                classes_from_typing = node.names
-
-                for cls in classes_from_typing:
-                    cls_name = cls.name
-                    if cls_name == "ContextManager":
-                        self.module_imports_info_dict["context_manager_from_typing"] = True
-                    elif cls_name == "AsyncContextManager":
-                        self.module_imports_info_dict["async_context_manager_from_typing"] = True
+            elif not python_2_support_required and path not in CONTEXTLIB_ALIAS_ALLOWLIST and node.module == "typing":
+                for imported_class in node.names:
+                    imported_class_name = imported_class.name
+                    if imported_class_name in CONTEXT_MANAGER_ALIASES:
+                        add_contextlib_alias_error(node, imported_class_name)
 
             self.generic_visit(node)
 
+        if not python_2_support_required and path not in CONTEXTLIB_ALIAS_ALLOWLIST:
+
+            def visit_Attribute(self, node: ast.Attribute) -> None:
+                if isinstance(node.value, ast.Name) and node.value.id == "typing" and node.attr in CONTEXT_MANAGER_ALIASES:
+                    add_contextlib_alias_error(node, node.attr)
+                self.generic_visit(node)
+
         def visit_AnnAssign(self, node: ast.AnnAssign) -> None:
             self.old_syntax_finder().visit(node.annotation)