Don't expand global variables in body of a function with constrained type variables (#9882)

See first test case for the scenario where this causes issues (the second test case is already handled elsewhere and just added for completeness). This is an old issue, but it is important to fix it now because recent changes to PEP 484 about re-export force people to use function aliases.

Co-authored-by: Ivan Levkivskyi <[email protected]>

Author: ilevkivskyi

mypy/test/testtransform.py

@@ -60,6 +60,7 @@ def test_transform(testcase: DataDrivenTestCase) -> None:
                     and not os.path.splitext(
                         os.path.basename(f.path))[0].endswith('_')):
                 t = TypeAssertTransformVisitor()
+                t.test_only = True
                 f = t.mypyfile(f)
                 a += str(f).split('\n')
     except CompileError as e:

mypy/treetransform.py

@@ -20,7 +20,7 @@
     YieldFromExpr, NamedTupleExpr, TypedDictExpr, NonlocalDecl, SetComprehension,
     DictionaryComprehension, ComplexExpr, TypeAliasExpr, EllipsisExpr,
     YieldExpr, ExecStmt, Argument, BackquoteExpr, AwaitExpr, AssignmentExpr,
-    OverloadPart, EnumCallExpr, REVEAL_TYPE
+    OverloadPart, EnumCallExpr, REVEAL_TYPE, GDEF
 )
 from mypy.types import Type, FunctionLike, ProperType
 from mypy.traverser import TraverserVisitor
@@ -37,6 +37,8 @@ class TransformVisitor(NodeVisitor[Node]):
 
     Notes:
 
+     * This can only be used to transform functions or classes, not top-level
+       statements, and/or modules as a whole.
      * Do not duplicate TypeInfo nodes. This would generally not be desirable.
      * Only update some name binding cross-references, but only those that
        refer to Var, Decorator or FuncDef nodes, not those targeting ClassDef or
@@ -48,6 +50,9 @@ class TransformVisitor(NodeVisitor[Node]):
     """
 
     def __init__(self) -> None:
+        # To simplify testing, set this flag to True if you want to transform
+        # all statements in a file (this is prohibited in normal mode).
+        self.test_only = False
         # There may be multiple references to a Var node. Keep track of
         # Var translations using a dictionary.
         self.var_map = {}  # type: Dict[Var, Var]
@@ -58,6 +63,7 @@ def __init__(self) -> None:
         self.func_placeholder_map = {}  # type: Dict[FuncDef, FuncDef]
 
     def visit_mypy_file(self, node: MypyFile) -> MypyFile:
+        assert self.test_only, "This visitor should not be used for whole files."
         # NOTE: The 'names' and 'imports' instance variables will be empty!
         ignored_lines = {line: codes[:]
                          for line, codes in node.ignored_lines.items()}
@@ -358,7 +364,10 @@ def copy_ref(self, new: RefExpr, original: RefExpr) -> None:
         new.fullname = original.fullname
         target = original.node
         if isinstance(target, Var):
-            target = self.visit_var(target)
+            # Do not transform references to global variables. See
+            # testGenericFunctionAliasExpand for an example where this is important.
+            if original.kind != GDEF:
+                target = self.visit_var(target)
         elif isinstance(target, Decorator):
             target = self.visit_var(target.var)
         elif isinstance(target, FuncDef):

test-data/unit/check-generics.test

@@ -2402,3 +2402,31 @@ def func(tp: Type[C[S]]) -> S:
     else:
         return reg[tp.test].test[0]
 [builtins fixtures/dict.pyi]
+
+[case testGenericFunctionAliasExpand]
+from typing import Optional, TypeVar
+
+T = TypeVar("T")
+def gen(x: T) -> T: ...
+gen_a = gen
+
+S = TypeVar("S", int, str)
+class C: ...
+def test() -> Optional[S]:
+    reveal_type(gen_a(C()))  # N: Revealed type is '__main__.C*'
+    return None
+
+[case testGenericFunctionMemberExpand]
+from typing import Optional, TypeVar, Callable
+
+T = TypeVar("T")
+
+class A:
+    def __init__(self) -> None:
+        self.gen: Callable[[T], T]
+
+S = TypeVar("S", int, str)
+class C: ...
+def test() -> Optional[S]:
+    reveal_type(A().gen(C()))  # N: Revealed type is '__main__.C*'
+    return None