Check binops with union operands (#531)

quora · May 3, 2022 · 416926b · 416926b
1 parent 9b3fc3c
commit 416926b
Show file tree

Hide file tree

Showing 4 changed files with 37 additions and 19 deletions.
diff --git a/docs/changelog.md b/docs/changelog.md
@@ -2,6 +2,7 @@
 
 ## Unreleased
 
+- Fix type checking of binary operators involving unions (#531)
 - Improve `TypeVar` solution heuristic for constrained
   typevars with multiple solutions (#532)
 - Fix resolution of stringified annotations in `__init__` methods (#530)

diff --git a/pyanalyze/name_check_visitor.py b/pyanalyze/name_check_visitor.py
@@ -3110,12 +3110,7 @@ def _visit_binop_internal(
                 )
             return value
 
-        (
-            description,
-            method,
-            imethod,
-            rmethod,
-        ) = BINARY_OPERATION_TO_DESCRIPTION_AND_METHOD[type(op)]
+        (_, method, imethod, _) = BINARY_OPERATION_TO_DESCRIPTION_AND_METHOD[type(op)]
         allow_call = allow_call and method not in self.options.get_value_for(
             DisallowCallsToDunders
         )
@@ -3133,14 +3128,31 @@ def _visit_binop_internal(
             if not inplace_errors:
                 return inplace_result
 
-        # TODO handle MVV properly here. The naive approach (removing this check)
-        # leads to an error on Union[int, float] + Union[int, float], presumably because
-        # some combinations need the left and some need the right variant.
-        # A proper solution may be to take the product of the MVVs on both sides and try
-        # them all.
-        if isinstance(left, MultiValuedValue) and isinstance(right, MultiValuedValue):
-            return AnyValue(AnySource.inference)
+        possibilities = []
+        for subval in flatten_values(left):
+            result = self._visit_binop_no_mvv(
+                Composite(subval, left_composite.varname, left_composite.node),
+                op,
+                right_composite,
+                source_node,
+                allow_call,
+            )
+            possibilities.append(result)
+        return unite_values(*possibilities)
 
+    def _visit_binop_no_mvv(
+        self,
+        left_composite: Composite,
+        op: ast.AST,
+        right_composite: Composite,
+        source_node: ast.AST,
+        allow_call: bool = True,
+    ) -> Value:
+        left = left_composite.value
+        right = right_composite.value
+        (description, method, _, rmethod) = BINARY_OPERATION_TO_DESCRIPTION_AND_METHOD[
+            type(op)
+        ]
         if rmethod is None:
             return self._check_dunder_call(
                 source_node,

diff --git a/pyanalyze/test_name_check_visitor.py b/pyanalyze/test_name_check_visitor.py
@@ -1129,10 +1129,7 @@ def capybara(ints: Sequence[Literal[1, 2]]):
             )
 
             lst3 = [i + j * 10 for i in range(2) for j in range(3)]
-            # TODO: should be list[int] instead
-            assert_is_value(
-                lst3, SequenceValue(list, [(True, AnyValue(AnySource.inference))])
-            )
+            assert_is_value(lst3, SequenceValue(list, [(True, TypedValue(int))]))
 
     @assert_passes()
     def test_dict_comprehension(self):

diff --git a/pyanalyze/test_operations.py b/pyanalyze/test_operations.py
@@ -124,8 +124,16 @@ def capybara(x):
     def test_union(self):
         from typing import Union
 
-        def capybara(x: Union[int, str]) -> None:
-            assert_is_value(x * 3, MultiValuedValue([TypedValue(int), TypedValue(str)]))
+        def capybara(
+            x: Union[int, str],
+            y: Union[int, str],
+            z: Union[int, float],
+            a: Union[int, float],
+        ) -> None:
+            assert_is_value(x * 3, TypedValue(int) | TypedValue(str))
+
+            x + y  # E: unsupported_operation
+            assert_is_value(z + a, TypedValue(int) | TypedValue(float))
 
     @assert_passes()
     def test_rop(self):