Repository: quora/pyanalyze
Pull request #500 (2022-03-12 00:04:28)

Improve type inference for calls to `set()`, `list()`, and `tuple()` with union arguments
Now list(Sequence[A] | Sequence[B]) will be list[A] | list[B] instead of list[A | B].



Patch:
diff --git a/docs/changelog.md b/docs/changelog.md
index f2bd637..97194f3 100644
--- a/docs/changelog.md
+++ b/docs/changelog.md
@@ -2,6 +2,8 @@
 
 ## Unreleased
 
+- Improve type inference for calls to `set()`, `list()`, and
+  `tuple()` with union arguments (#500)
 - Remove special-cased signatured for `sorted()` (#498)
 - Support type narrowing on `bool()` calls (#497)
 - Support context managers that may suppress exceptions (#496)
diff --git a/pyanalyze/implementation.py b/pyanalyze/implementation.py
index a1b7682..ab68023 100644
--- a/pyanalyze/implementation.py
+++ b/pyanalyze/implementation.py
@@ -315,35 +315,39 @@ def _super_impl(ctx: CallContext) -> Value:
         return KnownValue(super_val)
 
 
-def _tuple_impl(ctx: CallContext) -> Value:
+def _tuple_impl(ctx: CallContext) -> ImplReturn:
     return _sequence_impl(tuple, ctx)
 
 
-def _list_impl(ctx: CallContext) -> Value:
+def _list_impl(ctx: CallContext) -> ImplReturn:
     return _sequence_impl(list, ctx)
 
 
-def _set_impl(ctx: CallContext) -> Value:
+def _set_impl(ctx: CallContext) -> ImplReturn:
     return _sequence_impl(set, ctx)
 
 
-def _sequence_impl(typ: type, ctx: CallContext) -> Value:
+def _sequence_impl(typ: type, ctx: CallContext) -> ImplReturn:
     iterable = ctx.vars["iterable"]
     if iterable is _NO_ARG_SENTINEL:
-        return KnownValue(typ())
-    cvi = concrete_values_from_iterable(iterable, ctx.visitor)
-    if isinstance(cvi, CanAssignError):
-        ctx.show_error(
-            f"{iterable} is not iterable",
-            ErrorCode.unsupported_operation,
-            arg="iterable",
-            detail=str(cvi),
-        )
-        return TypedValue(typ)
-    elif isinstance(cvi, Value):
-        return GenericValue(typ, [cvi])
-    else:
-        return SequenceIncompleteValue.make_or_known(typ, cvi)
+        return ImplReturn(KnownValue(typ()))
+
+    def inner(iterable: Value) -> Value:
+        cvi = concrete_values_from_iterable(iterable, ctx.visitor)
+        if isinstance(cvi, CanAssignError):
+            ctx.show_error(
+                f"{iterable} is not iterable",
+                ErrorCode.unsupported_operation,
+                arg="iterable",
+                detail=str(cvi),
+            )
+            return TypedValue(typ)
+        elif isinstance(cvi, Value):
+            return GenericValue(typ, [cvi])
+        else:
+            return SequenceIncompleteValue.make_or_known(typ, cvi)
+
+    return flatten_unions(inner, iterable)
 
 
 def _list_append_impl(ctx: CallContext) -> ImplReturn:


Test patch:
diff --git a/pyanalyze/test_implementation.py b/pyanalyze/test_implementation.py
index b5a0508..6b48252 100644
--- a/pyanalyze/test_implementation.py
+++ b/pyanalyze/test_implementation.py
@@ -210,6 +210,19 @@ class TestSequenceImpl(TestNameCheckVisitorBase):
             # argument that is iterable but does not have __iter__
             assert_is_value(tuple(str(x)), GenericValue(tuple, [TypedValue(str)]))
 
+    @assert_passes()
+    def test_union(self):
+        from typing import Sequence, Union
+        from typing_extensions import Never
+
+        def capybara(x: Union[Sequence[int], Sequence[str]], never: Never):
+            assert_is_value(
+                tuple(x),
+                GenericValue(tuple, [TypedValue(int)])
+                | GenericValue(tuple, [TypedValue(str)]),
+            )
+            assert_is_value(tuple(never), GenericValue(tuple, [NO_RETURN_VALUE]))
+
     @assert_passes()
     def test_not_iterable(self):
         def capybara(x):