simplified away second branch

Author: randolf-scholz

mypy/argmap.py

@@ -4,15 +4,16 @@
 
 from collections.abc import Sequence
 from typing import TYPE_CHECKING, Callable
+from typing_extensions import TypeGuard
 
 from mypy import nodes
-from mypy.join import join_type_list
 from mypy.maptype import map_instance_to_supertype
 from mypy.typeops import make_simplified_union
 from mypy.types import (
     AnyType,
     Instance,
     ParamSpecType,
+    ProperType,
     TupleType,
     Type,
     TypedDictType,
@@ -62,10 +63,11 @@ def map_actuals_to_formals(
             if (
                 isinstance(actualt, UnionType)
                 and actualt.items
-                and is_equal_sized_tuples(actualt.items)
+                and is_equal_sized_tuples(
+                    proper_types := [get_proper_type(t) for t in actualt.items]
+                )
             ):
-                # Arbitrarily pick the first item in the union.
-                actualt = get_proper_type(actualt.items[0])
+                actualt = proper_types[0]
             if isinstance(actualt, TupleType):
                 # A tuple actual maps to a fixed number of formals.
                 for _ in range(len(actualt.items)):
@@ -215,18 +217,38 @@ def expand_actual_type(
         actual_type = get_proper_type(actual_type)
         if actual_kind == nodes.ARG_STAR:
             if isinstance(actual_type, UnionType):
-                # special case 1: union of equal sized tuples.  (e.g. `tuple[int, int] | tuple[None, None]`)
-                # special case 2: union contains no static sized tuples. (e.g. `list[str | None] | list[str]`)
-                if is_equal_sized_tuples(actual_type.items) or not any(
-                    isinstance(get_proper_type(t), TupleType) for t in actual_type.items
-                ):
-                    # If the actual type is a union, try expanding it.
-                    # Example:   f(*args), where args is `list[str | None] | list[str]`,
-                    # Example:   f(*args), where args is `tuple[A, B, C] | tuple[None, None, None]`
-                    # Note: there is potential for combinatorial explosion here:
-                    # f(*x1, *x2, .. *xn), if xₖ is a union of nₖ differently sized tuples,
-                    # then there are n₁ * n₂ * ... * nₖ possible combinations of pointer positions.
-                    # therefore, we only take this branch if all union members consume the same number of items.
+                proper_types = [get_proper_type(t) for t in actual_type.items]
+                # special case: union of equal sized tuples.  (e.g. `tuple[int, int] | tuple[None, None]`)
+                if is_equal_sized_tuples(proper_types):
+                    # transform union of tuples into a tuple of unions
+                    # e.g. tuple[A, B, C] | tuple[None, None, None] -> tuple[A | None, B | None, C | None]
+                    tuple_args: list[Type] = [
+                        make_simplified_union(items)
+                        for items in zip(*(t.items for t in proper_types))
+                    ]
+                    actual_type = TupleType(
+                        tuple_args,
+                        # use Iterable[A | B | C] as the fallback type
+                        fallback=Instance(
+                            self.context.iterable_type.type, [UnionType.make_union(tuple_args)]
+                        ),
+                    )
+                else:
+                    # reinterpret all union items as iterable types (if possible)
+                    # and return the union of the iterable item types results.
+                    from mypy.subtypes import is_subtype
+
+                    iterable_type = self.context.iterable_type
+
+                    def as_iterable_type(t: Type) -> Type:
+                        """Map a type to the iterable supertype if it is a subtype."""
+                        p_t = get_proper_type(t)
+                        if isinstance(p_t, Instance) and is_subtype(t, iterable_type):
+                            return map_instance_to_supertype(p_t, iterable_type.type)
+                        if isinstance(p_t, TupleType):
+                            # Convert tuple[A, B, C] to Iterable[A | B | C].
+                            return Instance(iterable_type.type, [make_simplified_union(p_t.items)])
+                        return t
 
                     # create copies of self for each item in the union
                     sub_expanders = [
@@ -239,7 +261,11 @@ def expand_actual_type(
                     candidate_type = make_simplified_union(
                         [
                             e.expand_actual_type(
-                                item, actual_kind, formal_name, formal_kind, allow_unpack
+                                as_iterable_type(item),
+                                actual_kind,
+                                formal_name,
+                                formal_kind,
+                                allow_unpack,
                             )
                             for e, item in zip(sub_expanders, actual_type.items)
                         ]
@@ -249,28 +275,6 @@ def expand_actual_type(
                     self.tuple_index = int(sub_expanders[0].tuple_index)
                     self.kwargs_used = set(sub_expanders[0].kwargs_used)
                     return candidate_type
-                else:
-                    # otherwise, we fall back to checking using the join of the union members.
-                    # for better results we first map all instances to Iterable[T]
-                    from mypy.subtypes import is_subtype
-
-                    iterable_type = self.context.iterable_type
-
-                    def as_iterable_type(t: Type) -> Type:
-                        """Map a type to the iterable supertype if it is a subtype."""
-                        p_t = get_proper_type(t)
-                        if isinstance(p_t, Instance) and is_subtype(t, iterable_type):
-                            return map_instance_to_supertype(p_t, iterable_type.type)
-                        if isinstance(p_t, TupleType):
-                            # Convert tuple[A, B, C] to Iterable[A | B | C].
-                            return Instance(iterable_type.type, [make_simplified_union(p_t.items)])
-                        return t
-
-                    joined_type = join_type_list([as_iterable_type(t) for t in actual_type.items])
-                    assert not isinstance(get_proper_type(joined_type), TupleType)
-                    return self.expand_actual_type(
-                        joined_type, actual_kind, formal_name, formal_kind, allow_unpack
-                    )
 
             if isinstance(actual_type, TypeVarTupleType):
                 # This code path is hit when *Ts is passed to a callable and various
@@ -346,19 +350,18 @@ def as_iterable_type(t: Type) -> Type:
             return original_actual
 
 
-def is_equal_sized_tuples(types: Sequence[Type]) -> bool:
+def is_equal_sized_tuples(types: Sequence[ProperType]) -> TypeGuard[Sequence[TupleType]]:
     """Check if all types are tuples of the same size."""
     if not types:
         return True
 
     iterator = iter(types)
-    first = get_proper_type(next(iterator))
+    first = next(iterator)
     if not isinstance(first, TupleType):
         return False
     size = first.length()
 
     for item in iterator:
-        p_t = get_proper_type(item)
-        if not isinstance(p_t, TupleType) or p_t.length() != size:
+        if not isinstance(item, TupleType) or item.length() != size:
             return False
     return True