Support generic partial types for attributes (python#8044)

ilevkivskyi · web-flow · commit 4e021d908a82 · 2019-12-03T14:00:36.000Z
Work towards python#1055 Currently partial types are supported for local variables. However, only partial `None` types are supported for `self` attributes. This PR adds the same level of support to generic partial types. They follow mostly the same rules: * A partial type can be refined in the _same_ method where it is defined. * But a partial type from class body can not be refined in a method, as if `local_partial_types = True`. The logic is pretty simple: the `.node` attribute for `self.attr` expressions is set to `None`, so I added a little helper to get it from the class symbol table instead.
diff --git a/mypy/checker.py b/mypy/checker.py
@@ -2153,13 +2153,20 @@ def try_infer_partial_generic_type_from_assignment(self,
             if foo():
                 x = [1]  # Infer List[int] as type of 'x'
         """
+        var = None
         if (isinstance(lvalue, NameExpr)
                 and isinstance(lvalue.node, Var)
                 and isinstance(lvalue.node.type, PartialType)):
             var = lvalue.node
-            typ = lvalue.node.type
+        elif isinstance(lvalue, MemberExpr):
+            var = self.expr_checker.get_partial_self_var(lvalue)
+        if var is not None:
+            typ = var.type
+            assert isinstance(typ, PartialType)
             if typ.type is None:
                 return
+            # TODO: some logic here duplicates the None partial type counterpart
+            #       inlined in check_assignment(), see # 8043.
             partial_types = self.find_partial_types(var)
             if partial_types is None:
                 return
@@ -2993,8 +3000,12 @@ def check_indexed_assignment(self, lvalue: IndexExpr,
     def try_infer_partial_type_from_indexed_assignment(
             self, lvalue: IndexExpr, rvalue: Expression) -> None:
         # TODO: Should we share some of this with try_infer_partial_type?
+        var = None
         if isinstance(lvalue.base, RefExpr) and isinstance(lvalue.base.node, Var):
             var = lvalue.base.node
+        elif isinstance(lvalue.base, MemberExpr):
+            var = self.expr_checker.get_partial_self_var(lvalue.base)
+        if isinstance(var, Var):
             if isinstance(var.type, PartialType):
                 type_type = var.type.type
                 if type_type is None:
@@ -4331,7 +4342,14 @@ def find_partial_types_in_all_scopes(
                 # All scopes within the outermost function are active. Scopes out of
                 # the outermost function are inactive to allow local reasoning (important
                 # for fine-grained incremental mode).
-                scope_active = (not self.options.local_partial_types
+                disallow_other_scopes = self.options.local_partial_types
+
+                if isinstance(var.type, PartialType) and var.type.type is not None and var.info:
+                    # This is an ugly hack to make partial generic self attributes behave
+                    # as if --local-partial-types is always on (because it used to be like this).
+                    disallow_other_scopes = True
+
+                scope_active = (not disallow_other_scopes
                                 or scope.is_local == self.partial_types[-1].is_local)
                 return scope_active, scope.is_local, scope.map
         return False, False, None
diff --git a/mypy/checkexpr.py b/mypy/checkexpr.py
@@ -537,6 +537,25 @@ def check_typeddict_call_with_kwargs(self, callee: TypedDictType,
 
         return callee
 
+    def get_partial_self_var(self, expr: MemberExpr) -> Optional[Var]:
+        """Get variable node for a partial self attribute.
+
+        If the expression is not a self attribute, or attribute is not variable,
+        or variable is not partial, return None.
+        """
+        if not (isinstance(expr.expr, NameExpr) and
+                isinstance(expr.expr.node, Var) and expr.expr.node.is_self):
+            # Not a self.attr expression.
+            return None
+        info = self.chk.scope.enclosing_class()
+        if not info or expr.name not in info.names:
+            # Don't mess with partial types in superclasses.
+            return None
+        sym = info.names[expr.name]
+        if isinstance(sym.node, Var) and isinstance(sym.node.type, PartialType):
+            return sym.node
+        return None
+
     # Types and methods that can be used to infer partial types.
     item_args = {'builtins.list': ['append'],
                  'builtins.set': ['add', 'discard'],
@@ -550,6 +569,8 @@ def check_typeddict_call_with_kwargs(self, callee: TypedDictType,
     def try_infer_partial_type(self, e: CallExpr) -> None:
         if isinstance(e.callee, MemberExpr) and isinstance(e.callee.expr, RefExpr):
             var = e.callee.expr.node
+            if var is None and isinstance(e.callee.expr, MemberExpr):
+                var = self.get_partial_self_var(e.callee.expr)
             if not isinstance(var, Var):
                 return
             partial_types = self.chk.find_partial_types(var)
diff --git a/test-data/unit/check-inference.test b/test-data/unit/check-inference.test
@@ -1458,9 +1458,9 @@ class A:
 class A:
     def f(self) -> None:
         # Attributes aren't supported right now.
-        self.a = [] # E: Need type annotation for 'a' (hint: "a: List[<type>] = ...")
+        self.a = []
         self.a.append(1)
-        self.a.append('')
+        self.a.append('')  # E: Argument 1 to "append" of "list" has incompatible type "str"; expected "int"
 [builtins fixtures/list.pyi]
 
 [case testInferListInitializedToEmptyInClassBodyAndOverriden]
@@ -1585,6 +1585,121 @@ oo.update(d)
 reveal_type(oo) # N: Revealed type is 'collections.OrderedDict[builtins.int*, builtins.str*]'
 [builtins fixtures/dict.pyi]
 
+[case testInferAttributeInitializedToEmptyAndAssigned]
+class C:
+    def __init__(self) -> None:
+        self.a = []
+        if bool():
+            self.a = [1]
+reveal_type(C().a)  # N: Revealed type is 'builtins.list[builtins.int*]'
+[builtins fixtures/list.pyi]
+
+[case testInferAttributeInitializedToEmptyAndAppended]
+class C:
+    def __init__(self) -> None:
+        self.a = []
+        if bool():
+            self.a.append(1)
+reveal_type(C().a)  # N: Revealed type is 'builtins.list[builtins.int]'
+[builtins fixtures/list.pyi]
+
+[case testInferAttributeInitializedToEmptyAndAssignedItem]
+class C:
+    def __init__(self) -> None:
+        self.a = {}
+        if bool():
+            self.a[0] = 'yes'
+reveal_type(C().a)  # N: Revealed type is 'builtins.dict[builtins.int, builtins.str]'
+[builtins fixtures/dict.pyi]
+
+[case testInferAttributeInitializedToNoneAndAssigned]
+# flags: --strict-optional
+class C:
+    def __init__(self) -> None:
+        self.a = None
+        if bool():
+            self.a = 1
+reveal_type(C().a)  # N: Revealed type is 'Union[builtins.int, None]'
+
+[case testInferAttributeInitializedToEmptyNonSelf]
+class C:
+    def __init__(self) -> None:
+        self.a = []  # E: Need type annotation for 'a' (hint: "a: List[<type>] = ...")
+        if bool():
+            a = self
+            a.a = [1]
+            a.a.append(1)
+reveal_type(C().a)  # N: Revealed type is 'builtins.list[Any]'
+[builtins fixtures/list.pyi]
+
+[case testInferAttributeInitializedToEmptyAndAssignedOtherMethod]
+class C:
+    def __init__(self) -> None:
+        self.a = []  # E: Need type annotation for 'a' (hint: "a: List[<type>] = ...")
+    def meth(self) -> None:
+        self.a = [1]
+reveal_type(C().a)  # N: Revealed type is 'builtins.list[Any]'
+[builtins fixtures/list.pyi]
+
+[case testInferAttributeInitializedToEmptyAndAppendedOtherMethod]
+class C:
+    def __init__(self) -> None:
+        self.a = []  # E: Need type annotation for 'a' (hint: "a: List[<type>] = ...")
+    def meth(self) -> None:
+        self.a.append(1)
+reveal_type(C().a)  # N: Revealed type is 'builtins.list[Any]'
+[builtins fixtures/list.pyi]
+
+[case testInferAttributeInitializedToEmptyAndAssignedItemOtherMethod]
+class C:
+    def __init__(self) -> None:
+        self.a = {}  # E: Need type annotation for 'a' (hint: "a: Dict[<type>, <type>] = ...")
+    def meth(self) -> None:
+        self.a[0] = 'yes'
+reveal_type(C().a)  # N: Revealed type is 'builtins.dict[Any, Any]'
+[builtins fixtures/dict.pyi]
+
+[case testInferAttributeInitializedToNoneAndAssignedOtherMethod]
+# flags: --strict-optional
+class C:
+    def __init__(self) -> None:
+        self.a = None
+    def meth(self) -> None:
+        self.a = 1  # E: Incompatible types in assignment (expression has type "int", variable has type "None")
+reveal_type(C().a)  # N: Revealed type is 'None'
+
+[case testInferAttributeInitializedToEmptyAndAssignedClassBody]
+class C:
+    a = []  # E: Need type annotation for 'a' (hint: "a: List[<type>] = ...")
+    def __init__(self) -> None:
+        self.a = [1]
+reveal_type(C().a)  # N: Revealed type is 'builtins.list[Any]'
+[builtins fixtures/list.pyi]
+
+[case testInferAttributeInitializedToEmptyAndAppendedClassBody]
+class C:
+    a = []  # E: Need type annotation for 'a' (hint: "a: List[<type>] = ...")
+    def __init__(self) -> None:
+        self.a.append(1)
+reveal_type(C().a)  # N: Revealed type is 'builtins.list[Any]'
+[builtins fixtures/list.pyi]
+
+[case testInferAttributeInitializedToEmptyAndAssignedItemClassBody]
+class C:
+    a = {}  # E: Need type annotation for 'a' (hint: "a: Dict[<type>, <type>] = ...")
+    def __init__(self) -> None:
+        self.a[0] = 'yes'
+reveal_type(C().a)  # N: Revealed type is 'builtins.dict[Any, Any]'
+[builtins fixtures/dict.pyi]
+
+[case testInferAttributeInitializedToNoneAndAssignedClassBody]
+# flags: --strict-optional
+class C:
+    a = None
+    def __init__(self) -> None:
+        self.a = 1
+reveal_type(C().a)  # N: Revealed type is 'Union[builtins.int, None]'
+
 
 -- Inferring types of variables first initialized to None (partial types)
 -- ----------------------------------------------------------------------
diff --git a/test-data/unit/fine-grained.test b/test-data/unit/fine-grained.test
@@ -2714,6 +2714,7 @@ class C:
         class D:
             def __init__(self) -> None:
                 self.x = {}
+            def meth(self) -> None:
                 self.x['a'] = 'b'
 [file a.py]
 def g() -> None: pass
@@ -2731,6 +2732,7 @@ class D:
     def __init__(self) -> None:
         a.g()
         self.x = {}
+    def meth(self) -> None:
         self.x['a'] = 'b'
 [file a.py]
 def g() -> None: pass
@@ -2742,6 +2744,69 @@ main:5: error: Need type annotation for 'x' (hint: "x: Dict[<type>, <type>] = ..
 ==
 main:5: error: Need type annotation for 'x' (hint: "x: Dict[<type>, <type>] = ...")
 
+[case testRefreshPartialTypeInferredAttributeIndex]
+from c import C
+reveal_type(C().a)
+[file c.py]
+from b import f
+class C:
+    def __init__(self) -> None:
+        self.a = {}
+        if bool():
+            self.a[0] = f()
+[file b.py]
+def f() -> int: ...
+[file b.py.2]
+from typing import List
+def f() -> str: ...
+[builtins fixtures/dict.pyi]
+[out]
+main:2: note: Revealed type is 'builtins.dict[builtins.int, builtins.int]'
+==
+main:2: note: Revealed type is 'builtins.dict[builtins.int, builtins.str]'
+
+[case testRefreshPartialTypeInferredAttributeAssign]
+from c import C
+reveal_type(C().a)
+[file c.py]
+from b import f
+class C:
+    def __init__(self) -> None:
+        self.a = []
+        if bool():
+            self.a = f()
+[file b.py]
+from typing import List
+def f() -> List[int]: ...
+[file b.py.2]
+from typing import List
+def f() -> List[str]: ...
+[builtins fixtures/list.pyi]
+[out]
+main:2: note: Revealed type is 'builtins.list[builtins.int]'
+==
+main:2: note: Revealed type is 'builtins.list[builtins.str]'
+
+[case testRefreshPartialTypeInferredAttributeAppend]
+from c import C
+reveal_type(C().a)
+[file c.py]
+from b import f
+class C:
+    def __init__(self) -> None:
+        self.a = []
+        if bool():
+            self.a.append(f())
+[file b.py]
+def f() -> int: ...
+[file b.py.2]
+def f() -> str: ...
+[builtins fixtures/list.pyi]
+[out]
+main:2: note: Revealed type is 'builtins.list[builtins.int]'
+==
+main:2: note: Revealed type is 'builtins.list[builtins.str]'
+
 [case testRefreshTryExcept]
 import a
 def f() -> None:
