Add support for operators with union operands

This pull request resolves python#2128 --
it modifies how we check operators to add support for operations like
`Union[int, float] + Union[int, float]`.

This approach basically iterates over all possible variations of the
left and right operands when they're unions and uses the union of the
resulting inferred type as the type of the overall expression.

Some implementation notes:

1.  I attempting "destructuring" just the left operand, which is
    basically the approach proposed here: python#2128 (comment)

    Unfortunately, I discovered it became necessary to also destructure
    the right operand to handle certain edge cases -- see the
    testOperatorDoubleUnionInterwovenUnionAdd test case.

2.  This algorithm varies slightly from what we do for union math
    in that we don't attempt to "preserve" the union/we always
    destructure both operands.

    I'm fairly confident that this is type-safe; I plan on testing
    this pull request against some internal code bases to help
    us gain more confidence.

Author: Michael0x2a

mypy/checker.py

@@ -2525,7 +2525,7 @@ def visit_operator_assignment_stmt(self,
         if inplace:
             # There is __ifoo__, treat as x = x.__ifoo__(y)
             rvalue_type, method_type = self.expr_checker.check_op(
-                method, lvalue_type, s.rvalue, s)
+                method, s.lvalue, lvalue_type, s.rvalue, s)
             if not is_subtype(rvalue_type, lvalue_type):
                 self.msg.incompatible_operator_assignment(s.op, s)
         else:

mypy/checkexpr.py

@@ -1697,7 +1697,7 @@ def visit_op_expr(self, e: OpExpr) -> Type:
 
         if e.op in nodes.op_methods:
             method = self.get_operator_method(e.op)
-            result, method_type = self.check_op(method, left_type, e.right, e,
+            result, method_type = self.check_op(method, e.left, left_type, e.right, e,
                                                 allow_reverse=True)
             e.method_type = method_type
             return result
@@ -1749,7 +1749,7 @@ def visit_comparison_expr(self, e: ComparisonExpr) -> Type:
                     sub_result = self.bool_type()
             elif operator in nodes.op_methods:
                 method = self.get_operator_method(operator)
-                sub_result, method_type = self.check_op(method, left_type, right, e,
+                sub_result, method_type = self.check_op(method, left, left_type, right, e,
                                                     allow_reverse=True)
 
             elif operator == 'is' or operator == 'is not':
@@ -1820,19 +1820,12 @@ def check_op_reversible(self,
                             left_expr: Expression,
                             right_type: Type,
                             right_expr: Expression,
-                            context: Context) -> Tuple[Type, Type]:
-        # Note: this kludge exists mostly to maintain compatibility with
-        # existing error messages. Apparently, if the left-hand-side is a
-        # union and we have a type mismatch, we print out a special,
-        # abbreviated error message. (See messages.unsupported_operand_types).
-        unions_present = isinstance(left_type, UnionType)
-
+                            context: Context,
+                            msg: MessageBuilder) -> Tuple[Type, Type]:
         def make_local_errors() -> MessageBuilder:
             """Creates a new MessageBuilder object."""
-            local_errors = self.msg.clean_copy()
+            local_errors = msg.clean_copy()
             local_errors.disable_count = 0
-            if unions_present:
-                local_errors.disable_type_names += 1
             return local_errors
 
         def lookup_operator(op_name: str, base_type: Type) -> Optional[Type]:
@@ -2009,9 +2002,9 @@ def lookup_definer(typ: Instance, attr_name: str) -> Optional[str]:
                 # TODO: Remove this extra case
                 return result
 
-        self.msg.add_errors(errors[0])
+        msg.add_errors(errors[0])
         if warn_about_uncalled_reverse_operator:
-            self.msg.reverse_operator_method_never_called(
+            msg.reverse_operator_method_never_called(
                 nodes.op_methods_to_symbols[op_name],
                 op_name,
                 right_type,
@@ -2025,22 +2018,57 @@ def lookup_definer(typ: Instance, attr_name: str) -> Optional[str]:
             result = error_any, error_any
             return result
 
-    def check_op(self, method: str, base_type: Type, arg: Expression,
-                 context: Context,
+    def check_op(self, method: str, base_expr: Expression, base_type: Type,
+                 arg: Expression, context: Context,
                  allow_reverse: bool = False) -> Tuple[Type, Type]:
         """Type check a binary operation which maps to a method call.
 
         Return tuple (result type, inferred operator method type).
         """
 
         if allow_reverse:
-            return self.check_op_reversible(
-                op_name=method,
-                left_type=base_type,
-                left_expr=TempNode(base_type),
-                right_type=self.accept(arg),
-                right_expr=arg,
-                context=context)
+            # Note: We want to pass in the original 'base_expr' and 'arg' for
+            # 'left_expr' and 'right_expr' whenever possible so that plugins
+            # and similar things can introspect on the original node if possible.
+            left_variants = [(base_type, base_expr)]
+            if isinstance(base_type, UnionType):
+                left_variants = [(item, TempNode(item)) for item in base_type.relevant_items()]
+
+            right_type = self.accept(arg)
+            right_variants = [(right_type, arg)]
+            if isinstance(right_type, UnionType):
+                right_variants = [(item, TempNode(item)) for item in right_type.relevant_items()]
+
+            msg = self.msg.clean_copy()
+            msg.disable_count = 0
+            all_results = []
+            all_inferred = []
+
+            for left_possible_type, left_expr in left_variants:
+                for right_possible_type, right_expr in right_variants:
+                    result, inferred = self.check_op_reversible(
+                        op_name=method,
+                        left_type=left_possible_type,
+                        left_expr=left_expr,
+                        right_type=right_possible_type,
+                        right_expr=right_expr,
+                        context=context,
+                        msg=msg)
+                    all_results.append(result)
+                    all_inferred.append(inferred)
+
+            if msg.is_errors():
+                self.msg.add_errors(msg)
+                if len(left_variants) >= 2 and len(right_variants) >= 2:
+                    self.msg.warn_both_operands_are_from_unions(context)
+                elif len(left_variants) >= 2:
+                    self.msg.warn_operand_was_from_union("Left", base_type, context)
+                elif len(right_variants) >= 2:
+                    self.msg.warn_operand_was_from_union("Right", right_type, context)
+
+            results_final = UnionType.make_simplified_union(all_results)
+            inferred_final = UnionType.make_simplified_union(all_inferred)
+            return results_final, inferred_final
         else:
             return self.check_op_local_by_name(
                 method=method,
@@ -2125,7 +2153,7 @@ def check_list_multiply(self, e: OpExpr) -> Type:
             left_type = self.accept(e.left, type_context=self.type_context[-1])
         else:
             left_type = self.accept(e.left)
-        result, method_type = self.check_op('__mul__', left_type, e.right, e)
+        result, method_type = self.check_op('__mul__', e.left, left_type, e.right, e)
         e.method_type = method_type
         return result
 
@@ -2179,7 +2207,7 @@ def visit_index_expr_helper(self, e: IndexExpr) -> Type:
               and left_type.is_type_obj() and left_type.type_object().is_enum):
             return self.visit_enum_index_expr(left_type.type_object(), e.index, e)
         else:
-            result, method_type = self.check_op('__getitem__', left_type, e.index, e)
+            result, method_type = self.check_op('__getitem__', e.base, left_type, e.index, e)
             e.method_type = method_type
             return result
 

mypy/messages.py

@@ -1019,6 +1019,12 @@ def reverse_operator_method_never_called(self,
             ),
             context=context)
 
+    def warn_both_operands_are_from_unions(self, context: Context) -> None:
+        self.note('Both left and right operands are unions', context)
+
+    def warn_operand_was_from_union(self, side: str, original: Type, context: Context) -> None:
+        self.note('{} operand is of type {}'.format(side, self.format(original)), context)
+
     def operator_method_signatures_overlap(
             self, reverse_class: TypeInfo, reverse_method: str, forward_class: Type,
             forward_method: str, context: Context) -> None:

test-data/unit/check-callable.test

@@ -46,7 +46,8 @@ from typing import Callable, Union
 x = 5  # type: Union[int, Callable[[], str], Callable[[], int]]
 
 if callable(x):
-    y = x() + 2 # E: Unsupported operand types for + (likely involving Union)
+    y = x() + 2 # E: Unsupported operand types for + ("str" and "int") \
+                # N: Left operand is of type "Union[str, int]"
 else:
     z = x + 6
 
@@ -60,7 +61,8 @@ x = 5  # type: Union[int, str, Callable[[], str]]
 if callable(x):
     y = x() + 'test'
 else:
-    z = x + 6  # E: Unsupported operand types for + (likely involving Union)
+    z = x + 6  # E: Unsupported operand types for + ("str" and "int") \
+               # N: Left operand is of type "Union[int, str]"
 
 [builtins fixtures/callable.pyi]
 
@@ -153,7 +155,8 @@ x = 5  # type: Union[int, Callable[[], str]]
 if callable(x) and x() == 'test':
     x()
 else:
-    x + 5  # E: Unsupported left operand type for + (some union)
+    x + 5  # E: Unsupported left operand type for + ("Callable[[], str]") \
+           # N: Left operand is of type "Union[int, Callable[[], str]]"
 
 [builtins fixtures/callable.pyi]
 

test-data/unit/check-classes.test

@@ -2019,6 +2019,136 @@ class FractionChild(Fraction): pass
 
 class A(metaclass=Real): pass
 
+[case testOperatorDoubleUnionIntFloat]
+from typing import Union
+
+a: Union[int, float]
+b: int
+c: float
+
+reveal_type(a + a)  # E: Revealed type is 'builtins.float'
+reveal_type(a + b)  # E: Revealed type is 'builtins.float'
+reveal_type(b + a)  # E: Revealed type is 'builtins.float'
+reveal_type(a + c)  # E: Revealed type is 'builtins.float'
+reveal_type(c + a)  # E: Revealed type is 'builtins.float'
+[builtins fixtures/ops.pyi]
+
+[case testOperatorDoubleUnionStandardSubtyping]
+from typing import Union
+
+class Parent:
+    def __add__(self, x: Parent) -> Parent: pass
+    def __radd__(self, x: Parent) -> Parent: pass
+
+class Child(Parent):
+    def __add__(self, x: Parent) -> Child: pass
+    def __radd__(self, x: Parent) -> Child: pass
+
+a: Union[Parent, Child]
+b: Parent
+c: Child
+
+reveal_type(a + a)  # E: Revealed type is '__main__.Parent'
+reveal_type(a + b)  # E: Revealed type is '__main__.Parent'
+reveal_type(b + a)  # E: Revealed type is '__main__.Parent'
+reveal_type(a + c)  # E: Revealed type is '__main__.Child'
+reveal_type(c + a)  # E: Revealed type is '__main__.Child'
+
+[case testOperatorDoubleUnionNoRelationship1]
+from typing import Union
+
+class Foo:
+    def __add__(self, x: Foo) -> Foo: pass
+    def __radd__(self, x: Foo) -> Foo: pass
+
+class Bar:
+    def __add__(self, x: Bar) -> Bar: pass
+    def __radd__(self, x: Bar) -> Bar: pass
+
+a: Union[Foo, Bar]
+b: Foo
+c: Bar
+
+a + a  # E: Unsupported operand types for + ("Foo" and "Bar") \
+       # E: Unsupported operand types for + ("Bar" and "Foo") \
+       # N: Both left and right operands are unions
+
+a + b  # E: Unsupported operand types for + ("Bar" and "Foo") \
+       # N: Left operand is of type "Union[Foo, Bar]"
+
+b + a  # E: Unsupported operand types for + ("Foo" and "Bar") \
+       # N: Right operand is of type "Union[Foo, Bar]"
+
+a + c  # E: Unsupported operand types for + ("Foo" and "Bar") \
+       # N: Left operand is of type "Union[Foo, Bar]"
+
+c + a  # E: Unsupported operand types for + ("Bar" and "Foo") \
+       # N: Right operand is of type "Union[Foo, Bar]"
+
+[case testOperatorDoubleUnionNoRelationship2]
+from typing import Union
+
+class Foo:
+    def __add__(self, x: Foo) -> Foo: pass
+    def __radd__(self, x: Foo) -> Foo: pass
+
+class Bar:
+    def __add__(self, x: Union[Foo, Bar]) -> Bar: pass
+    def __radd__(self, x: Union[Foo, Bar]) -> Bar: pass
+
+a: Union[Foo, Bar]
+b: Foo
+c: Bar
+
+reveal_type(a + a)  # E: Revealed type is 'Union[__main__.Foo, __main__.Bar]'
+reveal_type(a + b)  # E: Revealed type is 'Union[__main__.Foo, __main__.Bar]'
+reveal_type(b + a)  # E: Revealed type is 'Union[__main__.Foo, __main__.Bar]'
+reveal_type(a + c)  # E: Revealed type is '__main__.Bar'
+reveal_type(c + a)  # E: Revealed type is '__main__.Bar'
+
+[case testOperatorDoubleUnionNaiveAdd]
+from typing import Union
+
+class A: pass
+class B: pass
+class C:
+    def __radd__(self, x: A) -> int: pass
+class D:
+    def __radd__(self, x: B) -> str: pass
+
+x: Union[A, B]
+y: Union[C, D]
+
+x + y  # E: Unsupported operand types for + ("A" and "D") \
+       # E: Unsupported operand types for + ("B" and "C") \
+       # N: Both left and right operands are unions
+
+[case testOperatorDoubleUnionInterwovenUnionAdd]
+from typing import Union
+
+class Out1: pass
+class Out2: pass
+class Out3: pass
+class Out4: pass
+
+class A:
+    def __add__(self, x: D) -> Out1: pass
+class B:
+    def __add__(self, x: C) -> Out2: pass
+class C:
+    def __radd__(self, x: A) -> Out3: pass
+class D:
+    def __radd__(self, x: B) -> Out4: pass
+
+x: Union[A, B]
+y: Union[C, D]
+
+reveal_type(x + y)    # E: Revealed type is 'Union[__main__.Out3, __main__.Out1, __main__.Out2, __main__.Out4]'
+reveal_type(A() + y)  # E: Revealed type is 'Union[__main__.Out3, __main__.Out1]'
+reveal_type(B() + y)  # E: Revealed type is 'Union[__main__.Out2, __main__.Out4]'
+reveal_type(x + C())  # E: Revealed type is 'Union[__main__.Out3, __main__.Out2]'
+reveal_type(x + D())  # E: Revealed type is 'Union[__main__.Out1, __main__.Out4]'
+
 [case testAbstractReverseOperatorMethod]
 import typing
 from abc import abstractmethod

test-data/unit/check-generics.test

@@ -755,7 +755,8 @@ class Node(Generic[T]):
 UNode = Union[int, Node[T]]
 x = 1 # type: UNode[int]
 
-x + 1 # E: Unsupported left operand type for + (some union)
+x + 1 # E: Unsupported left operand type for + ("Node[int]") \
+      # N: Left operand is of type "Union[int, Node[int]]"
 if not isinstance(x, Node):
     x + 1
 

test-data/unit/check-incremental.test

@@ -3515,7 +3515,8 @@ from typing import Optional
 def foo() -> Optional[int]: return 0
 [out1]
 [out2]
-main:3: error: Unsupported operand types for + ("int" and "Optional[int]")
+main:3: error: Unsupported operand types for + ("int" and "None")
+main:3: note: Right operand is of type "Optional[int]"
 
 [case testAttrsIncrementalSubclassingCached]
 from a import A
@@ -4082,7 +4083,8 @@ class Baz:
         return 1
 [out]
 [out2]
-tmp/a.py:3: error: Unsupported operand types for + ("int" and "Optional[int]")
+tmp/a.py:3: error: Unsupported operand types for + ("int" and "None")
+tmp/a.py:3: note: Right operand is of type "Optional[int]"
 
 [case testIncrementalMetaclassUpdate]
 import a