Fix for NamedTuple in method [WIP] (#2553)

Fixes #2535.
python · Dec 14, 2016 · e193aed · e193aed
1 parent 0fe8670
commit e193aed
Show file tree

Hide file tree

Showing 4 changed files with 183 additions and 14 deletions.
diff --git a/mypy/semanal.py b/mypy/semanal.py
@@ -1681,14 +1681,18 @@ def check_namedtuple(self, node: Expression, var_name: str = None) -> Optional[T
         if not ok:
             # Error. Construct dummy return value.
             return self.build_namedtuple_typeinfo('namedtuple', [], [])
+        name = cast(StrExpr, call.args[0]).value
+        if name != var_name or self.is_func_scope():
+            # Give it a unique name derived from the line number.
+            name += '@' + str(call.line)
+        info = self.build_namedtuple_typeinfo(name, items, types)
+        # Store it as a global just in case it would remain anonymous.
+        # (Or in the nearest class if there is one.)
+        stnode = SymbolTableNode(GDEF, info, self.cur_mod_id)
+        if self.type:
+            self.type.names[name] = stnode
         else:
-            name = cast(StrExpr, call.args[0]).value
-            if name != var_name:
-                # Give it a unique name derived from the line number.
-                name += '@' + str(call.line)
-            info = self.build_namedtuple_typeinfo(name, items, types)
-            # Store it as a global just in case it would remain anonymous.
-            self.globals[name] = SymbolTableNode(GDEF, info, self.cur_mod_id)
+            self.globals[name] = stnode
         call.analyzed = NamedTupleExpr(info)
         call.analyzed.set_line(call.line, call.column)
         return info
@@ -1901,14 +1905,18 @@ def check_typeddict(self, node: Expression, var_name: str = None) -> Optional[Ty
         if not ok:
             # Error. Construct dummy return value.
             return self.build_typeddict_typeinfo('TypedDict', [], [])
+        name = cast(StrExpr, call.args[0]).value
+        if name != var_name or self.is_func_scope():
+            # Give it a unique name derived from the line number.
+            name += '@' + str(call.line)
+        info = self.build_typeddict_typeinfo(name, items, types)
+        # Store it as a global just in case it would remain anonymous.
+        # (Or in the nearest class if there is one.)
+        stnode = SymbolTableNode(GDEF, info, self.cur_mod_id)
+        if self.type:
+            self.type.names[name] = stnode
         else:
-            name = cast(StrExpr, call.args[0]).value
-            if name != var_name:
-                # Give it a unique name derived from the line number.
-                name += '@' + str(call.line)
-            info = self.build_typeddict_typeinfo(name, items, types)
-            # Store it as a global just in case it would remain anonymous.
-            self.globals[name] = SymbolTableNode(GDEF, info, self.cur_mod_id)
+            self.globals[name] = stnode
         call.analyzed = TypedDictExpr(info)
         call.analyzed.set_line(call.line, call.column)
         return info

diff --git a/test-data/unit/check-incremental.test b/test-data/unit/check-incremental.test
@@ -1640,3 +1640,129 @@ follow_imports = skip
 main:3: error: Revealed type is 'builtins.int'
 [out2]
 main:3: error: Revealed type is 'Any'
+
+[case testIncrementalNamedTupleInMethod]
+from ntcrash import nope
+[file ntcrash.py]
+from typing import NamedTuple
+class C:
+    def f(self) -> None:
+        A = NamedTuple('A', [('x', int), ('y', int)])
+[out1]
+main:1: error: Module 'ntcrash' has no attribute 'nope'
+[out2]
+main:1: error: Module 'ntcrash' has no attribute 'nope'
+
+[case testIncrementalNamedTupleInMethod2]
+from ntcrash import nope
+[file ntcrash.py]
+from typing import NamedTuple
+class C:
+    class D:
+        def f(self) -> None:
+            A = NamedTuple('A', [('x', int), ('y', int)])
+[out1]
+main:1: error: Module 'ntcrash' has no attribute 'nope'
+[out2]
+main:1: error: Module 'ntcrash' has no attribute 'nope'
+
+[case testIncrementalNamedTupleInMethod3]
+from ntcrash import nope
+[file ntcrash.py]
+from typing import NamedTuple
+class C:
+    def a(self):
+        class D:
+            def f(self) -> None:
+                A = NamedTuple('A', [('x', int), ('y', int)])
+[out1]
+main:1: error: Module 'ntcrash' has no attribute 'nope'
+[out2]
+main:1: error: Module 'ntcrash' has no attribute 'nope'
+
+[case testIncrementalNamedTupleInMethod4]
+from ntcrash import C
+reveal_type(C().a)
+reveal_type(C().b)
+reveal_type(C().c)
+[file ntcrash.py]
+from typing import NamedTuple
+class C:
+    def __init__(self) -> None:
+        A = NamedTuple('A', [('x', int)])
+        self.a = A(0)
+        self.b = A(0)  # type: A
+        self.c = A
+[out1]
+main:2: error: Revealed type is 'Tuple[builtins.int, fallback=ntcrash.C.A@4]'
+main:3: error: Revealed type is 'Tuple[builtins.int, fallback=ntcrash.C.A@4]'
+main:4: error: Revealed type is 'def (x: builtins.int) -> Tuple[builtins.int, fallback=ntcrash.C.A@4]'
+[out2]
+main:2: error: Revealed type is 'Tuple[builtins.int, fallback=ntcrash.C.A@4]'
+main:3: error: Revealed type is 'Tuple[builtins.int, fallback=ntcrash.C.A@4]'
+main:4: error: Revealed type is 'def (x: builtins.int) -> Tuple[builtins.int, fallback=ntcrash.C.A@4]'
+
+[case testIncrementalTypedDictInMethod]
+from tdcrash import nope
+[file tdcrash.py]
+from mypy_extensions import TypedDict
+class C:
+    def f(self) -> None:
+        A = TypedDict('A', {'x': int, 'y': int})
+[builtins fixtures/dict.pyi]
+[out1]
+main:1: error: Module 'tdcrash' has no attribute 'nope'
+[out2]
+main:1: error: Module 'tdcrash' has no attribute 'nope'
+
+[case testIncrementalTypedDictInMethod2]
+from tdcrash import nope
+[file tdcrash.py]
+from mypy_extensions import TypedDict
+class C:
+    class D:
+        def f(self) -> None:
+            A = TypedDict('A', {'x': int, 'y': int})
+[builtins fixtures/dict.pyi]
+[out1]
+main:1: error: Module 'tdcrash' has no attribute 'nope'
+[out2]
+main:1: error: Module 'tdcrash' has no attribute 'nope'
+
+[case testIncrementalTypedDictInMethod3]
+from tdcrash import nope
+[file tdcrash.py]
+from mypy_extensions import TypedDict
+class C:
+    def a(self):
+        class D:
+            def f(self) -> None:
+                A = TypedDict('A', {'x': int, 'y': int})
+[builtins fixtures/dict.pyi]
+[out1]
+main:1: error: Module 'tdcrash' has no attribute 'nope'
+[out2]
+main:1: error: Module 'tdcrash' has no attribute 'nope'
+
+[case testIncrementalTypedDictInMethod4]
+from ntcrash import C
+reveal_type(C().a)
+reveal_type(C().b)
+reveal_type(C().c)
+[file ntcrash.py]
+from mypy_extensions import TypedDict
+class C:
+    def __init__(self) -> None:
+        A = TypedDict('A', {'x': int})
+        self.a = A(x=0)
+        self.b = A(x=0)  # type: A
+        self.c = A
+[builtins fixtures/dict.pyi]
+[out1]
+main:2: error: Revealed type is 'TypedDict(x=builtins.int, _fallback=typing.Mapping[builtins.str, builtins.int])'
+main:3: error: Revealed type is 'TypedDict(x=builtins.int, _fallback=ntcrash.C.A@4)'
+main:4: error: Revealed type is 'def () -> ntcrash.C.A@4'
+[out2]
+main:2: error: Revealed type is 'TypedDict(x=builtins.int, _fallback=typing.Mapping[builtins.str, builtins.int])'
+main:3: error: Revealed type is 'TypedDict(x=builtins.int, _fallback=ntcrash.C.A@4)'
+main:4: error: Revealed type is 'def () -> ntcrash.C.A@4'
diff --git a/test-data/unit/check-namedtuple.test b/test-data/unit/check-namedtuple.test
@@ -411,3 +411,19 @@ reveal_type(B._make(['']))  # E: Revealed type is 'Tuple[builtins.str, fallback=
 b = B._make([''])  # type: B
 
 [builtins fixtures/list.pyi]
+
+[case testNamedTupleInClassNamespace]
+# https://github.com/python/mypy/pull/2553#issuecomment-266474341
+from typing import NamedTuple
+class C:
+    def f(self):
+        A = NamedTuple('A', [('x', int)])
+    def g(self):
+        A = NamedTuple('A', [('y', int)])
+C.A  # E: "C" has no attribute "A"
+
+[case testNamedTupleInFunction]
+from typing import NamedTuple
+def f() -> None:
+    A = NamedTuple('A', [('x', int)])
+A  # E: Name 'A' is not defined
diff --git a/test-data/unit/check-typeddict.test b/test-data/unit/check-typeddict.test
@@ -432,3 +432,22 @@ reveal_type(f(g))  # E: Revealed type is '<uninhabited>'
 
 -- TODO: Implement support for this case.
 --[case testCannotIsInstanceTypedDictType]
+
+-- scoping
+[case testTypedDictInClassNamespace]
+# https://github.com/python/mypy/pull/2553#issuecomment-266474341
+from mypy_extensions import TypedDict
+class C:
+    def f(self):
+        A = TypedDict('A', {'x': int})
+    def g(self):
+        A = TypedDict('A', {'y': int})
+C.A  # E: "C" has no attribute "A"
+[builtins fixtures/dict.pyi]
+
+[case testTypedDictInFunction]
+from mypy_extensions import TypedDict
+def f() -> None:
+    A = TypedDict('A', {'x': int})
+A  # E: Name 'A' is not defined
+[builtins fixtures/dict.pyi]