Fix binding of self-types on unions (#7152)

ilevkivskyi · web-flow · commit dab5b29fd6d5 · 2019-07-05T14:21:59.000+01:00
Fixes #7149 The fix is quite straightforward, self-types should be bound to each individual component, so I just added only basic tests. Please let me know if you want better test coverage (properties, class-methods) etc.
diff --git a/mypy/checkexpr.py b/mypy/checkexpr.py
@@ -677,7 +677,8 @@ def check_union_call_expr(self, e: CallExpr, object_type: UnionType, member: str
             self.msg.disable_errors()
             item = analyze_member_access(member, typ, e, False, False, False,
                                          self.msg, original_type=object_type, chk=self.chk,
-                                         in_literal_context=self.is_literal_context())
+                                         in_literal_context=self.is_literal_context(),
+                                         self_type=typ)
             self.msg.enable_errors()
             narrowed = self.narrow_type_from_binder(e.callee, item, skip_non_overlapping=True)
             if narrowed is None:
diff --git a/mypy/checkmember.py b/mypy/checkmember.py
@@ -45,11 +45,13 @@ def __init__(self,
                  original_type: Type,
                  context: Context,
                  msg: MessageBuilder,
-                 chk: 'mypy.checker.TypeChecker') -> None:
+                 chk: 'mypy.checker.TypeChecker',
+                 self_type: Optional[Type]) -> None:
         self.is_lvalue = is_lvalue
         self.is_super = is_super
         self.is_operator = is_operator
         self.original_type = original_type
+        self.self_type = self_type or original_type
         self.context = context  # Error context
         self.msg = msg
         self.chk = chk
@@ -60,9 +62,16 @@ def builtin_type(self, name: str) -> Instance:
     def not_ready_callback(self, name: str, context: Context) -> None:
         self.chk.handle_cannot_determine_type(name, context)
 
-    def copy_modified(self, messages: MessageBuilder) -> 'MemberContext':
-        return MemberContext(self.is_lvalue, self.is_super, self.is_operator,
-                             self.original_type, self.context, messages, self.chk)
+    def copy_modified(self, *, messages: Optional[MessageBuilder] = None,
+                      self_type: Optional[Type] = None) -> 'MemberContext':
+        mx = MemberContext(self.is_lvalue, self.is_super, self.is_operator,
+                           self.original_type, self.context, self.msg, self.chk,
+                           self.self_type)
+        if messages is not None:
+            mx.msg = messages
+        if self_type is not None:
+            mx.self_type = self_type
+        return mx
 
 
 def analyze_member_access(name: str,
@@ -75,7 +84,8 @@ def analyze_member_access(name: str,
                           original_type: Type,
                           chk: 'mypy.checker.TypeChecker',
                           override_info: Optional[TypeInfo] = None,
-                          in_literal_context: bool = False) -> Type:
+                          in_literal_context: bool = False,
+                          self_type: Optional[Type] = None) -> Type:
     """Return the type of attribute 'name' of 'typ'.
 
     The actual implementation is in '_analyze_member_access' and this docstring
@@ -91,15 +101,18 @@ def analyze_member_access(name: str,
     that we have available. When looking for an attribute of 'typ', we may perform
     recursive calls targeting the fallback type, and 'typ' may become some supertype
     of 'original_type'. 'original_type' is always preserved as the 'typ' type used in
-    the initial, non-recursive call.
+    the initial, non-recursive call. The 'self_type' is a component of 'original_type'
+    to which generic self should be bound (a narrower type that has a fallback to instance).
+    Currently this is used only for union types.
     """
     mx = MemberContext(is_lvalue,
                        is_super,
                        is_operator,
                        original_type,
                        context,
                        msg,
-                       chk=chk)
+                       chk=chk,
+                       self_type=self_type)
     result = _analyze_member_access(name, typ, mx, override_info)
     if in_literal_context and isinstance(result, Instance) and result.last_known_value is not None:
         return result.last_known_value
@@ -183,7 +196,7 @@ def analyze_instance_member_access(name: str,
             # the first argument.
             pass
         else:
-            signature = bind_self(signature, mx.original_type)
+            signature = bind_self(signature, mx.self_type)
         typ = map_instance_to_supertype(typ, method.info)
         member_type = expand_type_by_instance(signature, typ)
         freeze_type_vars(member_type)
@@ -263,8 +276,11 @@ def analyze_type_type_member_access(name: str, typ: TypeType, mx: MemberContext)
 
 def analyze_union_member_access(name: str, typ: UnionType, mx: MemberContext) -> Type:
     mx.msg.disable_type_names += 1
-    results = [_analyze_member_access(name, subtype, mx)
-               for subtype in typ.relevant_items()]
+    results = []
+    for subtype in typ.relevant_items():
+        # Self types should be bound to every individual item of a union.
+        item_mx = mx.copy_modified(self_type=subtype)
+        results.append(_analyze_member_access(name, subtype, item_mx))
     mx.msg.disable_type_names -= 1
     return UnionType.make_simplified_union(results)
 
@@ -342,7 +358,7 @@ def analyze_member_var_access(name: str,
                 # that the attribute exists
                 if method and method.info.fullname() != 'builtins.object':
                     function = function_type(method, mx.builtin_type('builtins.function'))
-                    bound_method = bind_self(function, mx.original_type)
+                    bound_method = bind_self(function, mx.self_type)
                     typ = map_instance_to_supertype(itype, method.info)
                     getattr_type = expand_type_by_instance(bound_method, typ)
                     if isinstance(getattr_type, CallableType):
@@ -359,7 +375,7 @@ def analyze_member_var_access(name: str,
             setattr_meth = info.get_method('__setattr__')
             if setattr_meth and setattr_meth.info.fullname() != 'builtins.object':
                 setattr_func = function_type(setattr_meth, mx.builtin_type('builtins.function'))
-                bound_type = bind_self(setattr_func, mx.original_type)
+                bound_type = bind_self(setattr_func, mx.self_type)
                 typ = map_instance_to_supertype(itype, setattr_meth.info)
                 setattr_type = expand_type_by_instance(bound_type, typ)
                 if isinstance(setattr_type, CallableType) and len(setattr_type.arg_types) > 0:
@@ -515,7 +531,7 @@ def analyze_var(name: str,
                 dispatched_type = meet.meet_types(mx.original_type, itype)
                 check_self_arg(functype, dispatched_type, var.is_classmethod, mx.context, name,
                                mx.msg)
-                signature = bind_self(functype, mx.original_type, var.is_classmethod)
+                signature = bind_self(functype, mx.self_type, var.is_classmethod)
                 if var.is_property:
                     # A property cannot have an overloaded type => the cast is fine.
                     assert isinstance(signature, CallableType)
diff --git a/test-data/unit/check-selftype.test b/test-data/unit/check-selftype.test
@@ -503,3 +503,33 @@ class C:
 ab: Union[A, B, C]
 reveal_type(ab.x)  # N: Revealed type is 'builtins.int'
 [builtins fixtures/property.pyi]
+
+[case testSelfTypeOnUnion]
+from typing import TypeVar, Union
+
+T = TypeVar('T')
+
+class A:
+    same: int
+
+class C:
+    def same(self: T) -> T: ...
+
+x: Union[A, C]
+reveal_type(x.same)  # N: Revealed type is 'Union[builtins.int, def () -> __main__.C*]'
+
+[case testSelfTypeOnUnionClassMethod]
+from typing import TypeVar, Union, Type
+
+T = TypeVar('T')
+
+class A:
+    same: int
+
+class C:
+    @classmethod
+    def same(cls: Type[T]) -> T: ...
+
+x: Union[A, C]
+reveal_type(x.same)  # N: Revealed type is 'Union[builtins.int, def () -> __main__.C*]'
+[builtins fixtures/classmethod.pyi]
diff --git a/test-data/unit/check-typeddict.test b/test-data/unit/check-typeddict.test
@@ -1662,3 +1662,14 @@ p = Point(x=42, y=1337)
 reveal_type(p)  # N: Revealed type is 'TypedDict('__main__.Point', {'x': builtins.int, 'y': builtins.int})'
 [builtins fixtures/dict.pyi]
 [typing fixtures/typing-full.pyi]
+
+[case testTypedDictOptionalUpdate]
+from typing import Union
+from mypy_extensions import TypedDict
+
+class A(TypedDict):
+    x: int
+
+d: Union[A, None]
+d.update({'x': 1})
+[builtins fixtures/dict.pyi]
