Repository: quora/pyanalyze
Pull request #518 (2022-04-12 23:48:04)

Fix propagation of no-return-unless constraints from calls to unions


Patch:
diff --git a/docs/changelog.md b/docs/changelog.md
index 805dd9a..91dfa7b 100644
--- a/docs/changelog.md
+++ b/docs/changelog.md
@@ -2,6 +2,8 @@
 
 ## Unreleased
 
+- Fix propagation of no-return-unless constraints from calls
+  to unions (#518)
 - Initial support for variable-length heterogeneous sequences
   (required for PEP 646). More precise types are now inferred
   for heterogeneous sequences containing variable-length
diff --git a/pyanalyze/name_check_visitor.py b/pyanalyze/name_check_visitor.py
index d91e287..f61ac94 100644
--- a/pyanalyze/name_check_visitor.py
+++ b/pyanalyze/name_check_visitor.py
@@ -4476,10 +4476,24 @@ class NameCheckVisitor(node_visitor.ReplacingNodeVisitor):
                     for val in callee.vals
                 ]
 
-            return unite_values(*values)
-        return self._check_call_no_mvv(
-            node, callee, args, keywords, allow_call=allow_call
-        )
+            pairs = [
+                unannotate_value(val, NoReturnConstraintExtension) for val in values
+            ]
+            val = unite_values(*[val for val, _ in pairs])
+            constraint = OrConstraint.make(
+                [
+                    AndConstraint.make(ext.constraint for ext in exts)
+                    for _, exts in pairs
+                ]
+            )
+        else:
+            val = self._check_call_no_mvv(
+                node, callee, args, keywords, allow_call=allow_call
+            )
+            val, nru_extensions = unannotate_value(val, NoReturnConstraintExtension)
+            constraint = AndConstraint.make(ext.constraint for ext in nru_extensions)
+        self.add_constraint(node, constraint)
+        return val
 
     def _check_call_no_mvv(
         self,
@@ -4528,12 +4542,6 @@ class NameCheckVisitor(node_visitor.ReplacingNodeVisitor):
                 with self.catch_errors():
                     return_value = extended_argspec.check_call(arguments, self, node)
 
-        return_value, nru_extensions = unannotate_value(
-            return_value, NoReturnConstraintExtension
-        )
-        for extension in nru_extensions:
-            self.add_constraint(node, extension.constraint)
-
         if extended_argspec is not None and not extended_argspec.has_return_value():
             local = self.get_local_return_value(extended_argspec)
             if local is not None:


Test patch:
diff --git a/pyanalyze/test_implementation.py b/pyanalyze/test_implementation.py
index b4f5f74..68af919 100644
--- a/pyanalyze/test_implementation.py
+++ b/pyanalyze/test_implementation.py
@@ -644,7 +644,6 @@ class TestGenericMutators(TestNameCheckVisitorBase):
             )
             lst.extend([3, 4])
 
-            # TODO: this is wrong; it drops all but the last Union member
             assert_is_value(
                 lst,
                 make_weak(
@@ -656,6 +655,16 @@ class TestGenericMutators(TestNameCheckVisitorBase):
                             )
                         ],
                     )
+                )
+                | make_weak(
+                    GenericValue(
+                        list,
+                        [
+                            MultiValuedValue(
+                                [KnownValue(1), KnownValue(3), KnownValue(4)]
+                            )
+                        ],
+                    )
                 ),
             )