rework starargs with union argument

Author: randolf-scholz

mypy/argmap.py

@@ -6,7 +6,9 @@
 from typing import TYPE_CHECKING, Callable
 
 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,
@@ -16,6 +18,7 @@
     TypedDictType,
     TypeOfAny,
     TypeVarTupleType,
+    UnionType,
     UnpackType,
     get_proper_type,
 )
@@ -54,6 +57,15 @@ def map_actuals_to_formals(
         elif actual_kind == nodes.ARG_STAR:
             # We need to know the actual type to map varargs.
             actualt = get_proper_type(actual_arg_type(ai))
+
+            # Special case for union of equal sized tuples.
+            if (
+                isinstance(actualt, UnionType)
+                and actualt.items
+                and is_equal_sized_tuples(actualt.items)
+            ):
+                # Arbitrarily pick the first item in the union.
+                actualt = get_proper_type(actualt.items[0])
             if isinstance(actualt, TupleType):
                 # A tuple actual maps to a fixed number of formals.
                 for _ in range(len(actualt.items)):
@@ -171,6 +183,15 @@ def __init__(self, context: ArgumentInferContext) -> None:
         # Type context for `*` and `**` arg kinds.
         self.context = context
 
+    def __eq__(self, other: object) -> bool:
+        if isinstance(other, ArgTypeExpander):
+            return (
+                self.tuple_index == other.tuple_index
+                and self.kwargs_used == other.kwargs_used
+                and self.context == other.context
+            )
+        return NotImplemented
+
     def expand_actual_type(
         self,
         actual_type: Type,
@@ -193,6 +214,64 @@ def expand_actual_type(
         original_actual = 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.
+
+                    # create copies of self for each item in the union
+                    sub_expanders = [
+                        ArgTypeExpander(context=self.context) for _ in actual_type.items
+                    ]
+                    for expander in sub_expanders:
+                        expander.tuple_index = int(self.tuple_index)
+                        expander.kwargs_used = set(self.kwargs_used)
+
+                    candidate_type = make_simplified_union(
+                        [
+                            e.expand_actual_type(
+                                item, actual_kind, formal_name, formal_kind, allow_unpack
+                            )
+                            for e, item in zip(sub_expanders, actual_type.items)
+                        ]
+                    )
+                    assert all(expander == sub_expanders[0] for expander in sub_expanders)
+                    # carry over the new state if all sub-expanders are the same state
+                    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
                 # special-handling didn't catch this. The best thing we can do is to use
@@ -265,3 +344,21 @@ def expand_actual_type(
         else:
             # No translation for other kinds -- 1:1 mapping.
             return original_actual
+
+
+def is_equal_sized_tuples(types: Sequence[Type]) -> bool:
+    """Check if all types are tuples of the same size."""
+    if not types:
+        return True
+
+    iterator = iter(types)
+    first = get_proper_type(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:
+            return False
+    return True

mypy/checkexpr.py

@@ -27,6 +27,7 @@
     freshen_function_type_vars,
 )
 from mypy.infer import ArgumentInferContext, infer_function_type_arguments, infer_type_arguments
+from mypy.join import join_type_list
 from mypy.literals import literal
 from mypy.maptype import map_instance_to_supertype
 from mypy.meet import is_overlapping_types, narrow_declared_type
@@ -5227,6 +5228,11 @@ def visit_tuple_expr(self, e: TupleExpr) -> Type:
                     ctx = None
                 tt = self.accept(item.expr, ctx)
                 tt = get_proper_type(tt)
+                if isinstance(tt, UnionType):
+                    # Coercing union to join allows better inference in some
+                    # special cases like `tuple[A, B] | tuple[C, D]`
+                    tt = get_proper_type(join_type_list(tt.items))
+
                 if isinstance(tt, TupleType):
                     if find_unpack_in_list(tt.items) is not None:
                         if seen_unpack_in_items:

mypy/constraints.py

@@ -3,24 +3,14 @@
 from __future__ import annotations
 
 from collections.abc import Iterable, Sequence
-from typing import TYPE_CHECKING, Final, cast
+from typing import Final, cast
 from typing_extensions import TypeGuard
 
 import mypy.subtypes
 import mypy.typeops
-from mypy.argmap import ArgTypeExpander
 from mypy.erasetype import erase_typevars
 from mypy.maptype import map_instance_to_supertype
-from mypy.nodes import (
-    ARG_OPT,
-    ARG_POS,
-    ARG_STAR,
-    ARG_STAR2,
-    CONTRAVARIANT,
-    COVARIANT,
-    ArgKind,
-    TypeInfo,
-)
+from mypy.nodes import ARG_OPT, ARG_POS, ARG_STAR, ARG_STAR2, CONTRAVARIANT, COVARIANT, TypeInfo
 from mypy.types import (
     TUPLE_LIKE_INSTANCE_NAMES,
     AnyType,
@@ -63,9 +53,6 @@
 from mypy.types_utils import is_union_with_any
 from mypy.typestate import type_state
 
-if TYPE_CHECKING:
-    from mypy.infer import ArgumentInferContext
-
 SUBTYPE_OF: Final = 0
 SUPERTYPE_OF: Final = 1
 
@@ -107,175 +94,6 @@ def __eq__(self, other: object) -> bool:
         return (self.type_var, self.op, self.target) == (other.type_var, other.op, other.target)
 
 
-def infer_constraints_for_callable(
-    callee: CallableType,
-    arg_types: Sequence[Type | None],
-    arg_kinds: list[ArgKind],
-    arg_names: Sequence[str | None] | None,
-    formal_to_actual: list[list[int]],
-    context: ArgumentInferContext,
-) -> list[Constraint]:
-    """Infer type variable constraints for a callable and actual arguments.
-
-    Return a list of constraints.
-    """
-    constraints: list[Constraint] = []
-    mapper = ArgTypeExpander(context)
-
-    param_spec = callee.param_spec()
-    param_spec_arg_types = []
-    param_spec_arg_names = []
-    param_spec_arg_kinds = []
-
-    incomplete_star_mapping = False
-    for i, actuals in enumerate(formal_to_actual):  # TODO: isn't this `enumerate(arg_types)`?
-        for actual in actuals:
-            if actual is None and callee.arg_kinds[i] in (ARG_STAR, ARG_STAR2):  # type: ignore[unreachable]
-                # We can't use arguments to infer ParamSpec constraint, if only some
-                # are present in the current inference pass.
-                incomplete_star_mapping = True  # type: ignore[unreachable]
-                break
-
-    for i, actuals in enumerate(formal_to_actual):
-        if isinstance(callee.arg_types[i], UnpackType):
-            unpack_type = callee.arg_types[i]
-            assert isinstance(unpack_type, UnpackType)
-
-            # In this case we are binding all the actuals to *args,
-            # and we want a constraint that the typevar tuple being unpacked
-            # is equal to a type list of all the actuals.
-            actual_types = []
-
-            unpacked_type = get_proper_type(unpack_type.type)
-            if isinstance(unpacked_type, TypeVarTupleType):
-                tuple_instance = unpacked_type.tuple_fallback
-            elif isinstance(unpacked_type, TupleType):
-                tuple_instance = unpacked_type.partial_fallback
-            else:
-                assert False, "mypy bug: unhandled constraint inference case"
-
-            for actual in actuals:
-                actual_arg_type = arg_types[actual]
-                if actual_arg_type is None:
-                    continue
-
-                expanded_actual = mapper.expand_actual_type(
-                    actual_arg_type,
-                    arg_kinds[actual],
-                    callee.arg_names[i],
-                    callee.arg_kinds[i],
-                    allow_unpack=True,
-                )
-
-                if arg_kinds[actual] != ARG_STAR or isinstance(
-                    get_proper_type(actual_arg_type), TupleType
-                ):
-                    actual_types.append(expanded_actual)
-                else:
-                    # If we are expanding an iterable inside * actual, append a homogeneous item instead
-                    actual_types.append(
-                        UnpackType(tuple_instance.copy_modified(args=[expanded_actual]))
-                    )
-
-            if isinstance(unpacked_type, TypeVarTupleType):
-                constraints.append(
-                    Constraint(
-                        unpacked_type,
-                        SUPERTYPE_OF,
-                        TupleType(actual_types, unpacked_type.tuple_fallback),
-                    )
-                )
-            elif isinstance(unpacked_type, TupleType):
-                # Prefixes get converted to positional args, so technically the only case we
-                # should have here is like Tuple[Unpack[Ts], Y1, Y2, Y3]. If this turns out
-                # not to hold we can always handle the prefixes too.
-                inner_unpack = unpacked_type.items[0]
-                assert isinstance(inner_unpack, UnpackType)
-                inner_unpacked_type = get_proper_type(inner_unpack.type)
-                suffix_len = len(unpacked_type.items) - 1
-                if isinstance(inner_unpacked_type, TypeVarTupleType):
-                    # Variadic item can be either *Ts...
-                    constraints.append(
-                        Constraint(
-                            inner_unpacked_type,
-                            SUPERTYPE_OF,
-                            TupleType(
-                                actual_types[:-suffix_len], inner_unpacked_type.tuple_fallback
-                            ),
-                        )
-                    )
-                else:
-                    # ...or it can be a homogeneous tuple.
-                    assert (
-                        isinstance(inner_unpacked_type, Instance)
-                        and inner_unpacked_type.type.fullname == "builtins.tuple"
-                    )
-                    for at in actual_types[:-suffix_len]:
-                        constraints.extend(
-                            infer_constraints(inner_unpacked_type.args[0], at, SUPERTYPE_OF)
-                        )
-                # Now handle the suffix (if any).
-                if suffix_len:
-                    for tt, at in zip(unpacked_type.items[1:], actual_types[-suffix_len:]):
-                        constraints.extend(infer_constraints(tt, at, SUPERTYPE_OF))
-            else:
-                assert False, "mypy bug: unhandled constraint inference case"
-        else:
-            for actual in actuals:
-                actual_arg_type = arg_types[actual]
-                if actual_arg_type is None:
-                    continue
-
-                if param_spec and callee.arg_kinds[i] in (ARG_STAR, ARG_STAR2):
-                    # If actual arguments are mapped to ParamSpec type, we can't infer individual
-                    # constraints, instead store them and infer single constraint at the end.
-                    # It is impossible to map actual kind to formal kind, so use some heuristic.
-                    # This inference is used as a fallback, so relying on heuristic should be OK.
-                    if not incomplete_star_mapping:
-                        param_spec_arg_types.append(
-                            mapper.expand_actual_type(
-                                actual_arg_type, arg_kinds[actual], None, arg_kinds[actual]
-                            )
-                        )
-                        actual_kind = arg_kinds[actual]
-                        param_spec_arg_kinds.append(
-                            ARG_POS if actual_kind not in (ARG_STAR, ARG_STAR2) else actual_kind
-                        )
-                        param_spec_arg_names.append(arg_names[actual] if arg_names else None)
-                else:
-                    actual_type = mapper.expand_actual_type(
-                        actual_arg_type,
-                        arg_kinds[actual],
-                        callee.arg_names[i],
-                        callee.arg_kinds[i],
-                    )
-                    c = infer_constraints(callee.arg_types[i], actual_type, SUPERTYPE_OF)
-                    constraints.extend(c)
-    if (
-        param_spec
-        and not any(c.type_var == param_spec.id for c in constraints)
-        and not incomplete_star_mapping
-    ):
-        # Use ParamSpec constraint from arguments only if there are no other constraints,
-        # since as explained above it is quite ad-hoc.
-        constraints.append(
-            Constraint(
-                param_spec,
-                SUPERTYPE_OF,
-                Parameters(
-                    arg_types=param_spec_arg_types,
-                    arg_kinds=param_spec_arg_kinds,
-                    arg_names=param_spec_arg_names,
-                    imprecise_arg_kinds=True,
-                ),
-            )
-        )
-    if any(isinstance(v, ParamSpecType) for v in callee.variables):
-        # As a perf optimization filter imprecise constraints only when we can have them.
-        constraints = filter_imprecise_kinds(constraints)
-    return constraints
-
-
 def infer_constraints(
     template: Type, actual: Type, direction: int, skip_neg_op: bool = False
 ) -> list[Constraint]: