From ae0f78d265e380a721efd27717e9f9189c655684 Mon Sep 17 00:00:00 2001
From: Ivan Levkivskyi <levkivskyi@gmail.com>
Date: Mon, 4 Mar 2019 00:29:40 +0000
Subject: [PATCH 01/13] Support plugin hooks on unions (normal methods)

---
 mypy/checkexpr.py                   | 52 +++++++++++++++++++++--------
 test-data/unit/check-typeddict.test | 22 ++++++++++++
 2 files changed, 61 insertions(+), 13 deletions(-)

diff --git a/mypy/checkexpr.py b/mypy/checkexpr.py
index d0687780c45c4..89e8ff3528ab4 100644
--- a/mypy/checkexpr.py
+++ b/mypy/checkexpr.py
@@ -297,6 +297,8 @@ def visit_call_expr_inner(self, e: CallExpr, allow_none_return: bool = False) ->
             return self.msg.untyped_function_call(callee_type, e)
         # Figure out the full name of the callee for plugin lookup.
         object_type = None
+        base_type = None
+        member = None
         if not isinstance(e.callee, RefExpr):
             fullname = None
         else:
@@ -306,15 +308,19 @@ def visit_call_expr_inner(self, e: CallExpr, allow_none_return: bool = False) ->
                 fullname = e.callee.node.target.type.fullname()
             if (fullname is None
                     and isinstance(e.callee, MemberExpr)
-                    and e.callee.expr in self.chk.type_map
-                    and isinstance(callee_type, FunctionLike)):
-                # For method calls we include the defining class for the method
-                # in the full name (example: 'typing.Mapping.get').
-                callee_expr_type = self.chk.type_map[e.callee.expr]
-                fullname = self.method_fullname(callee_expr_type, e.callee.name)
-                if fullname is not None:
-                    object_type = callee_expr_type
-        ret_type = self.check_call_expr_with_callee_type(callee_type, e, fullname, object_type)
+                    and e.callee.expr in self.chk.type_map):
+                if isinstance(callee_type, FunctionLike):
+                    # For method calls we include the defining class for the method
+                    # in the full name (example: 'typing.Mapping.get').
+                    callee_expr_type = self.chk.type_map[e.callee.expr]
+                    fullname = self.method_fullname(callee_expr_type, e.callee.name)
+                    if fullname is not None:
+                        object_type = callee_expr_type
+                if fullname is None:
+                    base_type = self.chk.type_map[e.callee.expr]
+                    member = e.callee.name
+        ret_type = self.check_call_expr_with_callee_type(callee_type, e, fullname,
+                                                         object_type, base_type, member)
         if isinstance(e.callee, RefExpr) and len(e.args) == 2:
             if e.callee.fullname in ('builtins.isinstance', 'builtins.issubclass'):
                 self.check_runtime_protocol_test(e)
@@ -645,15 +651,35 @@ def check_call_expr_with_callee_type(self,
                                          callee_type: Type,
                                          e: CallExpr,
                                          callable_name: Optional[str],
-                                         object_type: Optional[Type]) -> Type:
+                                         object_type: Optional[Type],
+                                         base_type: Optional[Type] = None,
+                                         member: Optional[str] = None) -> Type:
         """Type check call expression.
 
         The given callee type overrides the type of the callee
         expression.
         """
-        # Try to refine the call signature using plugin hooks before checking the call.
-        callee_type = self.transform_callee_type(
-            callable_name, callee_type, e.args, e.arg_kinds, e, e.arg_names, object_type)
+        if callable_name:
+            # Try to refine the call signature using plugin hooks before checking the call.
+            callee_type = self.transform_callee_type(
+                callable_name, callee_type, e.args, e.arg_kinds, e, e.arg_names, object_type)
+        elif isinstance(base_type, UnionType):
+            assert member
+            res = []  # type: List[Type]
+            for typ in base_type.items:
+                self.msg.disable_errors()
+                item = analyze_member_access(member, typ, e, False, False, False,
+                                             self.msg, original_type=base_type, chk=self.chk,
+                                             in_literal_context=self.is_literal_context())
+                self.msg.enable_errors()
+                item = self.narrow_type_from_binder(e.callee, item)
+                if isinstance(item, (NoneTyp, UninhabitedType)):
+                    continue
+                callable_name = self.method_fullname(typ, member)
+                object_type = typ if callable_name else None
+                res.append(self.check_call_expr_with_callee_type(item, e, callable_name,
+                                                                 object_type))
+            return UnionType.make_simplified_union(res)
 
         return self.check_call(callee_type, e.args, e.arg_kinds, e,
                                e.arg_names, callable_node=e.callee,
diff --git a/test-data/unit/check-typeddict.test b/test-data/unit/check-typeddict.test
index 92668c12bfd3c..619eda989b408 100644
--- a/test-data/unit/check-typeddict.test
+++ b/test-data/unit/check-typeddict.test
@@ -1571,3 +1571,25 @@ alias = b.__delitem__
 alias('x') # E: Argument 1 has incompatible type "str"; expected "NoReturn"
 alias(s) # E: Argument 1 has incompatible type "str"; expected "NoReturn"
 [builtins fixtures/dict.pyi]
+
+[case testPluginUnionsOfTypedDicts]
+from typing import Union
+from mypy_extensions import TypedDict
+
+class TDA(TypedDict):
+    a: int
+    b: str
+
+class TDB(TypedDict):
+    a: int
+    b: int
+    c: int
+
+td: Union[TDA, TDB]
+
+reveal_type(td.get('a'))  # E: Revealed type is 'builtins.int'
+reveal_type(td.get('b'))  # E: Revealed type is 'Union[builtins.str, builtins.int]'
+reveal_type(td.get('c'))  # E: Revealed type is 'Union[Any, builtins.int]' \
+                          # E: TypedDict "TDA" has no key 'c'
+[builtins fixtures/dict.pyi]
+[typing fixtures/typing-full.pyi]

From b3633bcde7afda47f4dd8b949be34b9c6130e54c Mon Sep 17 00:00:00 2001
From: Ivan Levkivskyi <levkivskyi@gmail.com>
Date: Mon, 4 Mar 2019 01:37:34 +0000
Subject: [PATCH 02/13] Support plugin hooks on unions (special methods)

---
 mypy/checkexpr.py                | 13 +++++++++++++
 test-data/unit/check-ctypes.test |  5 +++++
 2 files changed, 18 insertions(+)

diff --git a/mypy/checkexpr.py b/mypy/checkexpr.py
index 89e8ff3528ab4..00bf666f6c40d 100644
--- a/mypy/checkexpr.py
+++ b/mypy/checkexpr.py
@@ -2047,6 +2047,19 @@ def check_method_call_by_name(self,
         """
         local_errors = local_errors or self.msg
         original_type = original_type or base_type
+        if isinstance(base_type, UnionType):
+            res = []  # type: List[Type]
+            meth_res = []  # type: List[Type]
+            for typ in base_type.relevant_items():
+                local_errors.disable_type_names += 1
+                item, meth_item = self.check_method_call_by_name(method, typ, args, arg_kinds,
+                                                                 context, local_errors,
+                                                                 original_type)
+                local_errors.disable_type_names -= 1
+                res.append(item)
+                meth_res.append(meth_item)
+            return UnionType.make_simplified_union(res), UnionType.make_simplified_union(meth_res)
+
         method_type = analyze_member_access(method, base_type, context, False, False, True,
                                             local_errors, original_type=original_type,
                                             chk=self.chk,
diff --git a/test-data/unit/check-ctypes.test b/test-data/unit/check-ctypes.test
index 508b3bfd49c96..079db88e66a8f 100644
--- a/test-data/unit/check-ctypes.test
+++ b/test-data/unit/check-ctypes.test
@@ -23,6 +23,7 @@ for x in a:
 
 [case testCtypesArrayCustomElementType]
 import ctypes
+from typing import Union, List
 
 class MyCInt(ctypes.c_int):
     pass
@@ -46,6 +47,10 @@ mya[3] = b"bytes"  # E: No overload variant of "__setitem__" of "Array" matches
                    # N:     def __setitem__(self, slice, List[Union[MyCInt, int]]) -> None
 for myx in mya:
     reveal_type(myx)  # E: Revealed type is '__main__.MyCInt*'
+
+myu: Union[ctypes.Array[ctypes.c_int], List[str]]
+for myi in myu:
+    reveal_type(myi)  # E: Revealed type is 'Union[builtins.int*, builtins.str*]'
 [builtins fixtures/floatdict.pyi]
 
 [case testCtypesArrayUnionElementType]

From 15fafd3c74a6b45c77c750503414092bdb502b15 Mon Sep 17 00:00:00 2001
From: Ivan Levkivskyi <levkivskyi@gmail.com>
Date: Mon, 18 Mar 2019 15:39:47 +0000
Subject: [PATCH 03/13] Simplify logic

---
 mypy/checkexpr.py | 29 ++++++++++-------------------
 1 file changed, 10 insertions(+), 19 deletions(-)

diff --git a/mypy/checkexpr.py b/mypy/checkexpr.py
index 00bf666f6c40d..bf682de03ed25 100644
--- a/mypy/checkexpr.py
+++ b/mypy/checkexpr.py
@@ -297,7 +297,6 @@ def visit_call_expr_inner(self, e: CallExpr, allow_none_return: bool = False) ->
             return self.msg.untyped_function_call(callee_type, e)
         # Figure out the full name of the callee for plugin lookup.
         object_type = None
-        base_type = None
         member = None
         if not isinstance(e.callee, RefExpr):
             fullname = None
@@ -309,18 +308,10 @@ def visit_call_expr_inner(self, e: CallExpr, allow_none_return: bool = False) ->
             if (fullname is None
                     and isinstance(e.callee, MemberExpr)
                     and e.callee.expr in self.chk.type_map):
-                if isinstance(callee_type, FunctionLike):
-                    # For method calls we include the defining class for the method
-                    # in the full name (example: 'typing.Mapping.get').
-                    callee_expr_type = self.chk.type_map[e.callee.expr]
-                    fullname = self.method_fullname(callee_expr_type, e.callee.name)
-                    if fullname is not None:
-                        object_type = callee_expr_type
-                if fullname is None:
-                    base_type = self.chk.type_map[e.callee.expr]
-                    member = e.callee.name
+                member = e.callee.name
+                object_type = self.chk.type_map[e.callee.expr]
         ret_type = self.check_call_expr_with_callee_type(callee_type, e, fullname,
-                                                         object_type, base_type, member)
+                                                         object_type, member)
         if isinstance(e.callee, RefExpr) and len(e.args) == 2:
             if e.callee.fullname in ('builtins.isinstance', 'builtins.issubclass'):
                 self.check_runtime_protocol_test(e)
@@ -652,33 +643,33 @@ def check_call_expr_with_callee_type(self,
                                          e: CallExpr,
                                          callable_name: Optional[str],
                                          object_type: Optional[Type],
-                                         base_type: Optional[Type] = None,
                                          member: Optional[str] = None) -> Type:
         """Type check call expression.
 
         The given callee type overrides the type of the callee
         expression.
         """
+        if callable_name is None and member is not None:
+            callable_name = self.method_fullname(object_type, member)
         if callable_name:
             # Try to refine the call signature using plugin hooks before checking the call.
             callee_type = self.transform_callee_type(
                 callable_name, callee_type, e.args, e.arg_kinds, e, e.arg_names, object_type)
-        elif isinstance(base_type, UnionType):
-            assert member
+        elif member is not None and isinstance(object_type, UnionType):
             res = []  # type: List[Type]
-            for typ in base_type.items:
+            for typ in object_type.items:
                 self.msg.disable_errors()
                 item = analyze_member_access(member, typ, e, False, False, False,
-                                             self.msg, original_type=base_type, chk=self.chk,
+                                             self.msg, original_type=object_type, chk=self.chk,
                                              in_literal_context=self.is_literal_context())
                 self.msg.enable_errors()
                 item = self.narrow_type_from_binder(e.callee, item)
                 if isinstance(item, (NoneTyp, UninhabitedType)):
                     continue
                 callable_name = self.method_fullname(typ, member)
-                object_type = typ if callable_name else None
+                item_object_type = typ if callable_name else None
                 res.append(self.check_call_expr_with_callee_type(item, e, callable_name,
-                                                                 object_type))
+                                                                 item_object_type))
             return UnionType.make_simplified_union(res)
 
         return self.check_call(callee_type, e.args, e.arg_kinds, e,

From efd04a65a493697337a526f8e9a19277175b6d16 Mon Sep 17 00:00:00 2001
From: Ivan Levkivskyi <levkivskyi@gmail.com>
Date: Mon, 18 Mar 2019 16:06:48 +0000
Subject: [PATCH 04/13] Add comments and docstrings (and some more
 simplification)

---
 mypy/checkexpr.py | 30 ++++++++++++++++++++++++++----
 1 file changed, 26 insertions(+), 4 deletions(-)

diff --git a/mypy/checkexpr.py b/mypy/checkexpr.py
index bf682de03ed25..b67952ab64528 100644
--- a/mypy/checkexpr.py
+++ b/mypy/checkexpr.py
@@ -60,6 +60,7 @@
 from mypy.plugin import Plugin, MethodContext, MethodSigContext, FunctionContext
 from mypy.typeanal import make_optional_type
 from mypy.typeops import tuple_fallback
+from mypy import state
 
 # Type of callback user for checking individual function arguments. See
 # check_args() below for details.
@@ -298,13 +299,19 @@ def visit_call_expr_inner(self, e: CallExpr, allow_none_return: bool = False) ->
         # Figure out the full name of the callee for plugin lookup.
         object_type = None
         member = None
-        if not isinstance(e.callee, RefExpr):
-            fullname = None
-        else:
+        fullname = None
+        if isinstance(e.callee, RefExpr):
+            # There are two special cases where plugins might act:
+            # * A "static" reference/alias to a class or function;
+            #   get_function_hook() will be invoked for these.
             fullname = e.callee.fullname
             if (isinstance(e.callee.node, TypeAlias) and
                     isinstance(e.callee.node.target, Instance)):
                 fullname = e.callee.node.target.type.fullname()
+            # * Call to a method on object that has a full name (see
+            #   method_fullname() for details on supported objects);
+            #   get_method_hook() and get_method_signature_hook() will
+            #   be invoked for these.
             if (fullname is None
                     and isinstance(e.callee, MemberExpr)
                     and e.callee.expr in self.chk.type_map):
@@ -648,6 +655,11 @@ def check_call_expr_with_callee_type(self,
 
         The given callee type overrides the type of the callee
         expression.
+
+        The 'callable_name' and 'object_type' are used to call plugin hooks.
+        If 'callable_name' is None but 'member' is not None (member call), try constructing
+        'callable_name' using 'object_type' (the base type on which the method is called),
+        for example 'typing.Mapping.get'.
         """
         if callable_name is None and member is not None:
             callable_name = self.method_fullname(object_type, member)
@@ -655,16 +667,23 @@ def check_call_expr_with_callee_type(self,
             # Try to refine the call signature using plugin hooks before checking the call.
             callee_type = self.transform_callee_type(
                 callable_name, callee_type, e.args, e.arg_kinds, e, e.arg_names, object_type)
+        # Unions are special-cased to allow plugins to act on each item in the union.
         elif member is not None and isinstance(object_type, UnionType):
             res = []  # type: List[Type]
             for typ in object_type.items:
+                # Member access errors are already reported when visiting the member expression.
                 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())
                 self.msg.enable_errors()
                 item = self.narrow_type_from_binder(e.callee, item)
-                if isinstance(item, (NoneTyp, UninhabitedType)):
+                # TODO: This check is not one-to-one with non-special-cased behavior.
+                #       Ideas for more precise checks:
+                #       * reverse-engineering of mypy.meet.narrow_declared_type()
+                #       * apply special-casing only if we know some hooks will be applied
+                if (isinstance(item, NoneTyp) and not state.strict_optional or
+                        isinstance(item, UninhabitedType) and state.strict_optional):
                     continue
                 callable_name = self.method_fullname(typ, member)
                 item_object_type = typ if callable_name else None
@@ -2038,10 +2057,13 @@ def check_method_call_by_name(self,
         """
         local_errors = local_errors or self.msg
         original_type = original_type or base_type
+        # Unions are special-cased to allow plugins to act on each element of the union.
         if isinstance(base_type, UnionType):
             res = []  # type: List[Type]
             meth_res = []  # type: List[Type]
             for typ in base_type.relevant_items():
+                # Format error messages consistently with
+                # mypy.checkmember.analyze_union_member_access().
                 local_errors.disable_type_names += 1
                 item, meth_item = self.check_method_call_by_name(method, typ, args, arg_kinds,
                                                                  context, local_errors,

From 2dbbefbdf18a7202035085ed253281b263d2f362 Mon Sep 17 00:00:00 2001
From: Ivan Levkivskyi <levkivskyi@gmail.com>
Date: Mon, 18 Mar 2019 16:30:42 +0000
Subject: [PATCH 05/13] Add dedicated tests

---
 test-data/unit/check-custom-plugin.test | 44 +++++++++++++++++++++++++
 test-data/unit/plugins/union_method.py  | 40 ++++++++++++++++++++++
 2 files changed, 84 insertions(+)
 create mode 100644 test-data/unit/plugins/union_method.py

diff --git a/test-data/unit/check-custom-plugin.test b/test-data/unit/check-custom-plugin.test
index b810b8ba46b21..b90c108763a25 100644
--- a/test-data/unit/check-custom-plugin.test
+++ b/test-data/unit/check-custom-plugin.test
@@ -585,3 +585,47 @@ reveal_type(instance(2))  # E: Revealed type is 'builtins.float*'
 [file mypy.ini]
 [[mypy]
 plugins=<ROOT>/test-data/unit/plugins/callable_instance.py
+
+[case testGetMethodHooksOnUnions]
+# flags: --config-file tmp/mypy.ini --no-strict-optional
+from typing import Union
+
+class Foo:
+    def meth(self, x: str) -> str: ...
+class Bar:
+    def meth(self, x: int) -> str: ...
+class Other:
+    meth: int
+
+x: Union[Foo, Bar, Other]
+if isinstance(x.meth, int):
+    reveal_type(x.meth)  # E: Revealed type is 'builtins.int'
+else:
+    reveal_type(x.meth(int()))  # E: Revealed type is 'builtins.int'
+
+[builtins fixtures/isinstancelist.pyi]
+[file mypy.ini]
+[[mypy]
+plugins=<ROOT>/test-data/unit/plugins/union_method.py
+
+[case testGetMethodHooksOnUnionsStrictOptional]
+# flags: --config-file tmp/mypy.ini --strict-optional
+from typing import Union
+
+class Foo:
+    def meth(self, x: str) -> str: ...
+class Bar:
+    def meth(self, x: int) -> str: ...
+class Other:
+    meth: int
+
+x: Union[Foo, Bar, Other]
+if isinstance(x.meth, int):
+    reveal_type(x.meth)  # E: Revealed type is 'builtins.int'
+else:
+    reveal_type(x.meth(int()))  # E: Revealed type is 'builtins.int'
+
+[builtins fixtures/isinstancelist.pyi]
+[file mypy.ini]
+[[mypy]
+plugins=<ROOT>/test-data/unit/plugins/union_method.py
diff --git a/test-data/unit/plugins/union_method.py b/test-data/unit/plugins/union_method.py
new file mode 100644
index 0000000000000..d4dec03fe7c0b
--- /dev/null
+++ b/test-data/unit/plugins/union_method.py
@@ -0,0 +1,40 @@
+from mypy.plugin import (
+    CallableType, CheckerPluginInterface, MethodSigContext, MethodContext, Plugin
+)
+from mypy.types import Instance, Type
+
+
+class MethodPlugin(Plugin):
+    def get_method_signature_hook(self, fullname):
+        if fullname.startswith('__main__.Foo.'):
+            return my_meth_sig_hook
+        return None
+
+    def get_method_hook(self, fullname):
+        if fullname.startswith('__main__.Bar.'):
+            return my_meth_hook
+        return None
+
+
+def _str_to_int(api: CheckerPluginInterface, typ: Type) -> Type:
+    if isinstance(typ, Instance):
+        if typ.type.fullname() == 'builtins.str':
+            return api.named_generic_type('builtins.int', [])
+        elif typ.args:
+            return typ.copy_modified(args=[_str_to_int(api, t) for t in typ.args])
+    return typ
+
+
+def my_meth_sig_hook(ctx: MethodSigContext) -> CallableType:
+    return ctx.default_signature.copy_modified(
+        arg_types=[_str_to_int(ctx.api, t) for t in ctx.default_signature.arg_types],
+        ret_type=_str_to_int(ctx.api, ctx.default_signature.ret_type),
+    )
+
+
+def my_meth_hook(ctx: MethodContext) -> Type:
+    return _str_to_int(ctx.api, ctx.default_return_type)
+
+
+def plugin(version):
+    return MethodPlugin

From 72870a734220e34982cda1c70846aa916f287bc7 Mon Sep 17 00:00:00 2001
From: Ivan Levkivskyi <levkivskyi@gmail.com>
Date: Mon, 18 Mar 2019 16:41:50 +0000
Subject: [PATCH 06/13] Fix self-check

---
 mypy/checkexpr.py | 1 +
 1 file changed, 1 insertion(+)

diff --git a/mypy/checkexpr.py b/mypy/checkexpr.py
index b67952ab64528..4ad33d15bb227 100644
--- a/mypy/checkexpr.py
+++ b/mypy/checkexpr.py
@@ -662,6 +662,7 @@ def check_call_expr_with_callee_type(self,
         for example 'typing.Mapping.get'.
         """
         if callable_name is None and member is not None:
+            assert object_type is not None
             callable_name = self.method_fullname(object_type, member)
         if callable_name:
             # Try to refine the call signature using plugin hooks before checking the call.

From 8d1dcff02b59943cb50549a3e36d52dee6aa91a4 Mon Sep 17 00:00:00 2001
From: Ivan Levkivskyi <levkivskyi@gmail.com>
Date: Mon, 18 Mar 2019 16:54:49 +0000
Subject: [PATCH 07/13] More tests

---
 test-data/unit/check-custom-plugin.test | 17 ++++++++++++++++
 test-data/unit/check-typeddict.test     | 27 +++++++++++++++++++++++++
 2 files changed, 44 insertions(+)

diff --git a/test-data/unit/check-custom-plugin.test b/test-data/unit/check-custom-plugin.test
index b90c108763a25..30eb18b38cce4 100644
--- a/test-data/unit/check-custom-plugin.test
+++ b/test-data/unit/check-custom-plugin.test
@@ -629,3 +629,20 @@ else:
 [file mypy.ini]
 [[mypy]
 plugins=<ROOT>/test-data/unit/plugins/union_method.py
+
+[case testGetMethodHooksOnUnionsSpecial]
+# flags: --config-file tmp/mypy.ini
+from typing import Union
+
+class Foo:
+    def __getitem__(self, x: str) -> str: ...
+class Bar:
+    def __getitem__(self, x: int) -> str: ...
+
+x: Union[Foo, Bar]
+reveal_type(x[int()])  # E: Revealed type is 'builtins.int'
+
+[builtins fixtures/isinstancelist.pyi]
+[file mypy.ini]
+[[mypy]
+plugins=<ROOT>/test-data/unit/plugins/union_method.py
diff --git a/test-data/unit/check-typeddict.test b/test-data/unit/check-typeddict.test
index 619eda989b408..ba417d20ad2c2 100644
--- a/test-data/unit/check-typeddict.test
+++ b/test-data/unit/check-typeddict.test
@@ -1591,5 +1591,32 @@ reveal_type(td.get('a'))  # E: Revealed type is 'builtins.int'
 reveal_type(td.get('b'))  # E: Revealed type is 'Union[builtins.str, builtins.int]'
 reveal_type(td.get('c'))  # E: Revealed type is 'Union[Any, builtins.int]' \
                           # E: TypedDict "TDA" has no key 'c'
+
+reveal_type(td['a'])  # E: Revealed type is 'builtins.int'
+reveal_type(td['b'])  # E: Revealed type is 'Union[builtins.str, builtins.int]'
+reveal_type(td['c'])  # E: Revealed type is 'Union[Any, builtins.int]' \
+                      # E: TypedDict "TDA" has no key 'c'
+[builtins fixtures/dict.pyi]
+[typing fixtures/typing-full.pyi]
+
+[case testPluginUnionsOfTypedDictsNonTotal]
+from typing import Union
+from mypy_extensions import TypedDict
+
+class TDA(TypedDict, total=False):
+    a: int
+    b: str
+
+class TDB(TypedDict, total=False):
+    a: int
+    b: int
+    c: int
+
+td: Union[TDA, TDB]
+
+reveal_type(td.pop('a'))  # E: Revealed type is 'builtins.int'
+reveal_type(td.pop('b'))  # E: Revealed type is 'Union[builtins.str, builtins.int]'
+reveal_type(td.pop('c'))  # E: Revealed type is 'Union[Any, builtins.int]' \
+                          # E: TypedDict "TDA" has no key 'c'
 [builtins fixtures/dict.pyi]
 [typing fixtures/typing-full.pyi]

From 4d7aa10f1277255584e5b11fe7a247306756998d Mon Sep 17 00:00:00 2001
From: Ivan Levkivskyi <levkivskyi@gmail.com>
Date: Sat, 29 Jun 2019 20:32:42 +0100
Subject: [PATCH 08/13] Update after merge

---
 mypy/checkexpr.py                       |  2 +-
 test-data/unit/check-custom-plugin.test | 10 +++++-----
 2 files changed, 6 insertions(+), 6 deletions(-)

diff --git a/mypy/checkexpr.py b/mypy/checkexpr.py
index e6cea662c2773..412b4e9a4ad8c 100644
--- a/mypy/checkexpr.py
+++ b/mypy/checkexpr.py
@@ -670,7 +670,7 @@ def check_call_expr_with_callee_type(self,
                 #       Ideas for more precise checks:
                 #       * reverse-engineering of mypy.meet.narrow_declared_type()
                 #       * apply special-casing only if we know some hooks will be applied
-                if (isinstance(item, NoneTyp) and not state.strict_optional or
+                if (isinstance(item, NoneType) and not state.strict_optional or
                         isinstance(item, UninhabitedType) and state.strict_optional):
                     continue
                 callable_name = self.method_fullname(typ, member)
diff --git a/test-data/unit/check-custom-plugin.test b/test-data/unit/check-custom-plugin.test
index d762c9c55a5ec..93f9b1cecfa12 100644
--- a/test-data/unit/check-custom-plugin.test
+++ b/test-data/unit/check-custom-plugin.test
@@ -598,9 +598,9 @@ class Other:
 
 x: Union[Foo, Bar, Other]
 if isinstance(x.meth, int):
-    reveal_type(x.meth)  # E: Revealed type is 'builtins.int'
+    reveal_type(x.meth)  # N: Revealed type is 'builtins.int'
 else:
-    reveal_type(x.meth(int()))  # E: Revealed type is 'builtins.int'
+    reveal_type(x.meth(int()))  # N: Revealed type is 'builtins.int'
 
 [builtins fixtures/isinstancelist.pyi]
 [file mypy.ini]
@@ -620,9 +620,9 @@ class Other:
 
 x: Union[Foo, Bar, Other]
 if isinstance(x.meth, int):
-    reveal_type(x.meth)  # E: Revealed type is 'builtins.int'
+    reveal_type(x.meth)  # N: Revealed type is 'builtins.int'
 else:
-    reveal_type(x.meth(int()))  # E: Revealed type is 'builtins.int'
+    reveal_type(x.meth(int()))  # N: Revealed type is 'builtins.int'
 
 [builtins fixtures/isinstancelist.pyi]
 [file mypy.ini]
@@ -639,7 +639,7 @@ class Bar:
     def __getitem__(self, x: int) -> str: ...
 
 x: Union[Foo, Bar]
-reveal_type(x[int()])  # E: Revealed type is 'builtins.int'
+reveal_type(x[int()])  # N: Revealed type is 'builtins.int'
 
 [builtins fixtures/isinstancelist.pyi]
 [file mypy.ini]

From a643bbd47cff914ff7f9263dc47c3053f56375b9 Mon Sep 17 00:00:00 2001
From: Ivan Levkivskyi <levkivskyi@gmail.com>
Date: Sat, 29 Jun 2019 21:17:58 +0100
Subject: [PATCH 09/13] Address CR

---
 mypy/checkexpr.py | 99 +++++++++++++++++++++++++++++------------------
 1 file changed, 61 insertions(+), 38 deletions(-)

diff --git a/mypy/checkexpr.py b/mypy/checkexpr.py
index 412b4e9a4ad8c..c09fd188fde96 100644
--- a/mypy/checkexpr.py
+++ b/mypy/checkexpr.py
@@ -640,10 +640,11 @@ def check_call_expr_with_callee_type(self,
                                          member: Optional[str] = None) -> Type:
         """Type check call expression.
 
-        The given callee type overrides the type of the callee
-        expression.
+        The callee_type should be used as the type of callee expression. In particular,
+        in case of a union type this can be a particular item of the union, so that we can
+        apply plugin hooks to each item.
 
-        The 'callable_name' and 'object_type' are used to call plugin hooks.
+        The 'member', 'callable_name' and 'object_type' are only used to call plugin hooks.
         If 'callable_name' is None but 'member' is not None (member call), try constructing
         'callable_name' using 'object_type' (the base type on which the method is called),
         for example 'typing.Mapping.get'.
@@ -657,33 +658,36 @@ def check_call_expr_with_callee_type(self,
                 callable_name, callee_type, e.args, e.arg_kinds, e, e.arg_names, object_type)
         # Unions are special-cased to allow plugins to act on each item in the union.
         elif member is not None and isinstance(object_type, UnionType):
-            res = []  # type: List[Type]
-            for typ in object_type.items:
-                # Member access errors are already reported when visiting the member expression.
-                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())
-                self.msg.enable_errors()
-                item = self.narrow_type_from_binder(e.callee, item)
-                # TODO: This check is not one-to-one with non-special-cased behavior.
-                #       Ideas for more precise checks:
-                #       * reverse-engineering of mypy.meet.narrow_declared_type()
-                #       * apply special-casing only if we know some hooks will be applied
-                if (isinstance(item, NoneType) and not state.strict_optional or
-                        isinstance(item, UninhabitedType) and state.strict_optional):
-                    continue
-                callable_name = self.method_fullname(typ, member)
-                item_object_type = typ if callable_name else None
-                res.append(self.check_call_expr_with_callee_type(item, e, callable_name,
-                                                                 item_object_type))
-            return UnionType.make_simplified_union(res)
-
+            return self.check_union_call_expr(e, object_type, member)
         return self.check_call(callee_type, e.args, e.arg_kinds, e,
                                e.arg_names, callable_node=e.callee,
                                callable_name=callable_name,
                                object_type=object_type)[0]
 
+    def check_union_call_expr(self, e: CallExpr, object_type: UnionType, member: str) -> Type:
+        """"Type check calling a member expression where the base type is a union."""
+        res = []  # type: List[Type]
+        for typ in object_type.relevant_items():
+            # Member access errors are already reported when visiting the member expression.
+            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())
+            self.msg.enable_errors()
+            item = self.narrow_type_from_binder(e.callee, item)
+            # TODO: This check is not one-to-one with non-special-cased behavior.
+            #       Ideas for more precise checks:
+            #       * reverse-engineering of mypy.meet.narrow_declared_type()
+            #       * apply special-casing only if we know some hooks will be applied
+            if (isinstance(item, NoneType) and not state.strict_optional or
+                    isinstance(item, UninhabitedType) and state.strict_optional):
+                continue
+            callable_name = self.method_fullname(typ, member)
+            item_object_type = typ if callable_name else None
+            res.append(self.check_call_expr_with_callee_type(item, e, callable_name,
+                                                             item_object_type))
+        return UnionType.make_simplified_union(res)
+
     def check_call(self,
                    callee: Type,
                    args: List[Expression],
@@ -2057,19 +2061,9 @@ def check_method_call_by_name(self,
         original_type = original_type or base_type
         # Unions are special-cased to allow plugins to act on each element of the union.
         if isinstance(base_type, UnionType):
-            res = []  # type: List[Type]
-            meth_res = []  # type: List[Type]
-            for typ in base_type.relevant_items():
-                # Format error messages consistently with
-                # mypy.checkmember.analyze_union_member_access().
-                local_errors.disable_type_names += 1
-                item, meth_item = self.check_method_call_by_name(method, typ, args, arg_kinds,
-                                                                 context, local_errors,
-                                                                 original_type)
-                local_errors.disable_type_names -= 1
-                res.append(item)
-                meth_res.append(meth_item)
-            return UnionType.make_simplified_union(res), UnionType.make_simplified_union(meth_res)
+            return self.check_union_method_call_by_name(method, base_type,
+                                                        args, arg_kinds,
+                                                        context, local_errors, original_type)
 
         method_type = analyze_member_access(method, base_type, context, False, False, True,
                                             local_errors, original_type=original_type,
@@ -2078,6 +2072,35 @@ def check_method_call_by_name(self,
         return self.check_method_call(
             method, base_type, method_type, args, arg_kinds, context, local_errors)
 
+    def check_union_method_call_by_name(self,
+                                        method: str,
+                                        base_type: UnionType,
+                                        args: List[Expression],
+                                        arg_kinds: List[int],
+                                        context: Context,
+                                        local_errors: Optional[MessageBuilder] = None,
+                                        original_type: Optional[Type] = None
+                                        ) -> Tuple[Type, Type]:
+        """Type check a call to a named method on an object with union type.
+
+        This essentially checks the call using check_method_call_by_name() for each
+        union item and unions the result. We do this to allow plugins to act on
+        individual union items.
+        """
+        res = []  # type: List[Type]
+        meth_res = []  # type: List[Type]
+        for typ in base_type.relevant_items():
+            # Format error messages consistently with
+            # mypy.checkmember.analyze_union_member_access().
+            local_errors.disable_type_names += 1
+            item, meth_item = self.check_method_call_by_name(method, typ, args, arg_kinds,
+                                                             context, local_errors,
+                                                             original_type)
+            local_errors.disable_type_names -= 1
+            res.append(item)
+            meth_res.append(meth_item)
+        return UnionType.make_simplified_union(res), UnionType.make_simplified_union(meth_res)
+
     def check_method_call(self,
                           method_name: str,
                           base_type: Type,

From b3b8d89d5c316ea5655723caac639fcfb956cefb Mon Sep 17 00:00:00 2001
From: Ivan Levkivskyi <levkivskyi@gmail.com>
Date: Sat, 29 Jun 2019 21:57:10 +0100
Subject: [PATCH 10/13] A cleaner solution for non-overlapping types

---
 mypy/checkexpr.py | 34 +++++++++++++++++++++++-----------
 1 file changed, 23 insertions(+), 11 deletions(-)

diff --git a/mypy/checkexpr.py b/mypy/checkexpr.py
index c09fd188fde96..b69523655c5f7 100644
--- a/mypy/checkexpr.py
+++ b/mypy/checkexpr.py
@@ -3,7 +3,7 @@
 from collections import OrderedDict
 from contextlib import contextmanager
 from typing import (
-    cast, Dict, Set, List, Tuple, Callable, Union, Optional, Sequence, Iterator
+    cast, Dict, Set, List, Tuple, Callable, Union, Optional, Sequence, Iterator, overload
 )
 from typing_extensions import ClassVar, Final
 
@@ -674,17 +674,12 @@ def check_union_call_expr(self, e: CallExpr, object_type: UnionType, member: str
                                          self.msg, original_type=object_type, chk=self.chk,
                                          in_literal_context=self.is_literal_context())
             self.msg.enable_errors()
-            item = self.narrow_type_from_binder(e.callee, item)
-            # TODO: This check is not one-to-one with non-special-cased behavior.
-            #       Ideas for more precise checks:
-            #       * reverse-engineering of mypy.meet.narrow_declared_type()
-            #       * apply special-casing only if we know some hooks will be applied
-            if (isinstance(item, NoneType) and not state.strict_optional or
-                    isinstance(item, UninhabitedType) and state.strict_optional):
+            narrowed = self.narrow_type_from_binder(e.callee, item, skip_non_overlapping=True)
+            if narrowed is None:
                 continue
             callable_name = self.method_fullname(typ, member)
             item_object_type = typ if callable_name else None
-            res.append(self.check_call_expr_with_callee_type(item, e, callable_name,
+            res.append(self.check_call_expr_with_callee_type(narrowed, e, callable_name,
                                                              item_object_type))
         return UnionType.make_simplified_union(res)
 
@@ -2078,7 +2073,7 @@ def check_union_method_call_by_name(self,
                                         args: List[Expression],
                                         arg_kinds: List[int],
                                         context: Context,
-                                        local_errors: Optional[MessageBuilder] = None,
+                                        local_errors: MessageBuilder,
                                         original_type: Optional[Type] = None
                                         ) -> Tuple[Type, Type]:
         """Type check a call to a named method on an object with union type.
@@ -3600,7 +3595,19 @@ def bool_type(self) -> Instance:
         """Return instance type 'bool'."""
         return self.named_type('builtins.bool')
 
-    def narrow_type_from_binder(self, expr: Expression, known_type: Type) -> Type:
+    @overload
+    def narrow_type_from_binder(self, expr: Expression, known_type: Type) -> Type: ...
+    @overload
+    def narrow_type_from_binder(self, expr: Expression, known_type: Type,
+                                skip_non_overlapping: bool) -> Optional[Type]: ...
+
+    def narrow_type_from_binder(self, expr: Expression, known_type: Type,
+                                skip_non_overlapping: bool = False) -> Optional[Type]:
+        """Narrow down a known type of expression using information in conditional type binder.
+
+        If 'skip_non_overlapping' is True, return None if the type and restriction are
+        non-overlapping.
+        """
         if literal(expr) >= LITERAL_TYPE:
             restriction = self.chk.binder.get(expr)
             # If the current node is deferred, some variables may get Any types that they
@@ -3608,6 +3615,11 @@ def narrow_type_from_binder(self, expr: Expression, known_type: Type) -> Type:
             # produce invalid inferred Optional[Any] types, at least.
             if restriction and not (isinstance(known_type, AnyType)
                                     and self.chk.current_node_deferred):
+                # Note: this call should match the one in narrow_declared_type().
+                if (skip_non_overlapping and
+                        not is_overlapping_types(known_type, restriction,
+                                                 prohibit_none_typevar_overlap=True)):
+                    return None
                 ans = narrow_declared_type(known_type, restriction)
                 return ans
         return known_type

From 55b46242cb0411ad56fb4d461d15afe7e327cefc Mon Sep 17 00:00:00 2001
From: Ivan Levkivskyi <levkivskyi@gmail.com>
Date: Sat, 29 Jun 2019 22:00:56 +0100
Subject: [PATCH 11/13] Fix lint

---
 mypy/checkexpr.py | 6 +++---
 1 file changed, 3 insertions(+), 3 deletions(-)

diff --git a/mypy/checkexpr.py b/mypy/checkexpr.py
index b69523655c5f7..f434fa3a546ce 100644
--- a/mypy/checkexpr.py
+++ b/mypy/checkexpr.py
@@ -56,7 +56,6 @@
 from mypy.plugin import Plugin, MethodContext, MethodSigContext, FunctionContext
 from mypy.typeanal import make_optional_type
 from mypy.typeops import tuple_fallback
-from mypy import state
 
 # Type of callback user for checking individual function arguments. See
 # check_args() below for details.
@@ -3597,11 +3596,12 @@ def bool_type(self) -> Instance:
 
     @overload
     def narrow_type_from_binder(self, expr: Expression, known_type: Type) -> Type: ...
-    @overload
+
+    @overload  # noqa
     def narrow_type_from_binder(self, expr: Expression, known_type: Type,
                                 skip_non_overlapping: bool) -> Optional[Type]: ...
 
-    def narrow_type_from_binder(self, expr: Expression, known_type: Type,
+    def narrow_type_from_binder(self, expr: Expression, known_type: Type,  # noqa
                                 skip_non_overlapping: bool = False) -> Optional[Type]:
         """Narrow down a known type of expression using information in conditional type binder.
 

From 2bdb8bef76e3b859797f7b36716a56f23ef68644 Mon Sep 17 00:00:00 2001
From: Ivan Levkivskyi <levkivskyi@gmail.com>
Date: Sat, 29 Jun 2019 22:54:59 +0100
Subject: [PATCH 12/13] Try fixing Python 3.5.1

---
 mypy/checkexpr.py | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/mypy/checkexpr.py b/mypy/checkexpr.py
index f434fa3a546ce..eab335a479332 100644
--- a/mypy/checkexpr.py
+++ b/mypy/checkexpr.py
@@ -3,9 +3,9 @@
 from collections import OrderedDict
 from contextlib import contextmanager
 from typing import (
-    cast, Dict, Set, List, Tuple, Callable, Union, Optional, Sequence, Iterator, overload
+    cast, Dict, Set, List, Tuple, Callable, Union, Optional, Sequence, Iterator
 )
-from typing_extensions import ClassVar, Final
+from typing_extensions import ClassVar, Final, overload
 
 from mypy.errors import report_internal_error
 from mypy.typeanal import (

From d36ea3cb80bc1b4c06289a3e4c337c73346dae05 Mon Sep 17 00:00:00 2001
From: Ivan Levkivskyi <ilevkivskyi@dropbox.com>
Date: Thu, 4 Jul 2019 00:59:35 +0100
Subject: [PATCH 13/13] Address CR

---
 test-data/unit/check-custom-plugin.test |  6 +++---
 test-data/unit/plugins/union_method.py  | 11 ++++++++++-
 2 files changed, 13 insertions(+), 4 deletions(-)

diff --git a/test-data/unit/check-custom-plugin.test b/test-data/unit/check-custom-plugin.test
index 93f9b1cecfa12..ce3b1c4e665ea 100644
--- a/test-data/unit/check-custom-plugin.test
+++ b/test-data/unit/check-custom-plugin.test
@@ -592,7 +592,7 @@ from typing import Union
 class Foo:
     def meth(self, x: str) -> str: ...
 class Bar:
-    def meth(self, x: int) -> str: ...
+    def meth(self, x: int) -> float: ...
 class Other:
     meth: int
 
@@ -614,7 +614,7 @@ from typing import Union
 class Foo:
     def meth(self, x: str) -> str: ...
 class Bar:
-    def meth(self, x: int) -> str: ...
+    def meth(self, x: int) -> float: ...
 class Other:
     meth: int
 
@@ -636,7 +636,7 @@ from typing import Union
 class Foo:
     def __getitem__(self, x: str) -> str: ...
 class Bar:
-    def __getitem__(self, x: int) -> str: ...
+    def __getitem__(self, x: int) -> float: ...
 
 x: Union[Foo, Bar]
 reveal_type(x[int()])  # N: Revealed type is 'builtins.int'
diff --git a/test-data/unit/plugins/union_method.py b/test-data/unit/plugins/union_method.py
index d4dec03fe7c0b..28b9611ab9b06 100644
--- a/test-data/unit/plugins/union_method.py
+++ b/test-data/unit/plugins/union_method.py
@@ -25,6 +25,15 @@ def _str_to_int(api: CheckerPluginInterface, typ: Type) -> Type:
     return typ
 
 
+def _float_to_int(api: CheckerPluginInterface, typ: Type) -> Type:
+    if isinstance(typ, Instance):
+        if typ.type.fullname() == 'builtins.float':
+            return api.named_generic_type('builtins.int', [])
+        elif typ.args:
+            return typ.copy_modified(args=[_float_to_int(api, t) for t in typ.args])
+    return typ
+
+
 def my_meth_sig_hook(ctx: MethodSigContext) -> CallableType:
     return ctx.default_signature.copy_modified(
         arg_types=[_str_to_int(ctx.api, t) for t in ctx.default_signature.arg_types],
@@ -33,7 +42,7 @@ def my_meth_sig_hook(ctx: MethodSigContext) -> CallableType:
 
 
 def my_meth_hook(ctx: MethodContext) -> Type:
-    return _str_to_int(ctx.api, ctx.default_return_type)
+    return _float_to_int(ctx.api, ctx.default_return_type)
 
 
 def plugin(version):