diff --git a/mypy/expandtype.py b/mypy/expandtype.py index fd507216a6be9..bea0c81113d27 100644 --- a/mypy/expandtype.py +++ b/mypy/expandtype.py @@ -400,7 +400,57 @@ def expand_unpack(self, t: UnpackType) -> list[Type]: raise RuntimeError(f"Invalid type replacement to expand: {repl}") def visit_parameters(self, t: Parameters) -> Type: - return t.copy_modified(arg_types=self.expand_types(t.arg_types)) + arg_types: list[Type] = [] + arg_kinds: list[ArgKind] = [] + arg_names: list[str | None] = [] + for arg_type, arg_kind, arg_name in zip(t.arg_types, t.arg_kinds, t.arg_names): + expanded_arg_type: Type | None = None + expanded_vararg: list[Type] | None = None + tuple_fallback: Instance | None = None + if arg_kind == ARG_STAR and isinstance(arg_type, UnpackType): + if isinstance(arg_type.type, TypeVarTupleType): + expanded_vararg = self.expand_unpack(arg_type) + tuple_fallback = arg_type.type.tuple_fallback + else: + expanded_arg_type = arg_type.accept(self) + if isinstance(expanded_arg_type, UnpackType): + unpacked = get_proper_type(expanded_arg_type.type) + if isinstance(unpacked, TupleType): + expanded_vararg = unpacked.items + tuple_fallback = unpacked.partial_fallback + if expanded_vararg is not None: + # Keep a residual unpack and its suffix together as one vararg. Otherwise + # the suffix would become positional arguments placed after *args. + unpack_index = next( + (i for i, item in enumerate(expanded_vararg) if isinstance(item, UnpackType)), + None, + ) + if unpack_index is not None: + arg_types.extend(expanded_vararg[:unpack_index]) + arg_kinds.extend([ArgKind.ARG_POS] * unpack_index) + arg_names.extend([None] * unpack_index) + + unpack = expanded_vararg[unpack_index] + assert isinstance(unpack, UnpackType) + if unpack_index < len(expanded_vararg) - 1: + assert tuple_fallback is not None + unpack = UnpackType( + TupleType(expanded_vararg[unpack_index:], tuple_fallback) + ) + arg_types.append(unpack) + arg_kinds.append(ARG_STAR) + arg_names.append(arg_name) + else: + arg_types.extend(expanded_vararg) + arg_kinds.extend([ArgKind.ARG_POS] * len(expanded_vararg)) + arg_names.extend([None] * len(expanded_vararg)) + else: + arg_types.append( + expanded_arg_type if expanded_arg_type is not None else arg_type.accept(self) + ) + arg_kinds.append(arg_kind) + arg_names.append(arg_name) + return t.copy_modified(arg_types=arg_types, arg_kinds=arg_kinds, arg_names=arg_names) def interpolate_args_for_unpack(self, t: CallableType, var_arg: UnpackType) -> list[Type]: star_index = t.arg_kinds.index(ARG_STAR) diff --git a/mypy/test/testtypes.py b/mypy/test/testtypes.py index b287e82b3d4af..b9670712c26ad 100644 --- a/mypy/test/testtypes.py +++ b/mypy/test/testtypes.py @@ -40,6 +40,7 @@ LiteralType, NoneType, Overloaded, + Parameters, ProperType, SentinelValue, TupleType, @@ -300,6 +301,41 @@ def test_expand_naked_type_var(self) -> None: def test_expand_basic_generic_types(self) -> None: self.assert_expand(self.fx.gt, [(self.fx.t.id, self.fx.a)], self.fx.ga) + def test_expand_parameters_type_var_tuple_twice(self) -> None: + initial = Parameters( + [UnpackType(self.fx.ts), self.fx.d], + [ARG_STAR, ARG_NAMED], + ["args", "flag"], + variables=[self.fx.ts], + ) + first = mypy.expandtype.expand_type( + initial, + { + self.fx.ts.id: TupleType( + [self.fx.a, UnpackType(self.fx.us), self.fx.b], self.fx.std_tuple + ) + }, + ) + assert isinstance(first, Parameters) + assert first == Parameters( + [ + self.fx.a, + UnpackType(TupleType([UnpackType(self.fx.us), self.fx.b], self.fx.std_tuple)), + self.fx.d, + ], + [ARG_POS, ARG_STAR, ARG_NAMED], + [None, "args", "flag"], + ) + + second = mypy.expandtype.expand_type( + first, {self.fx.us.id: TupleType([self.fx.c], self.fx.std_tuple)} + ) + assert second == Parameters( + [self.fx.a, self.fx.c, self.fx.b, self.fx.d], + [ARG_POS, ARG_POS, ARG_POS, ARG_NAMED], + [None, None, None, "flag"], + ) + # IDEA: Add test cases for # tuple types # callable types @@ -1655,7 +1691,7 @@ def make_call(*items: tuple[str, str | None]) -> CallExpr: class TestExpandTypeLimitGetProperType(TestCase): # WARNING: do not increase this number unless absolutely necessary, # and you understand what you are doing. - ALLOWED_GET_PROPER_TYPES = 7 + ALLOWED_GET_PROPER_TYPES = 8 @skipUnless(mypy.expandtype.__file__.endswith(".py"), "Skip for compiled mypy") def test_count_get_proper_type(self) -> None: diff --git a/test-data/unit/check-parameter-specification.test b/test-data/unit/check-parameter-specification.test index f5eb4c416cda0..0e6425e808fe4 100644 --- a/test-data/unit/check-parameter-specification.test +++ b/test-data/unit/check-parameter-specification.test @@ -141,6 +141,38 @@ reveal_type(whatever) # N: Revealed type is "def (x: builtins.int) -> builtins. reveal_type(whatever(217)) # N: Revealed type is "builtins.list[builtins.int]" [builtins fixtures/paramspec.pyi] +[case testParamSpecVariadicContextManager] +from typing import Callable, Generic, TypeVar, TypeVarTuple, Unpack +from typing_extensions import ParamSpec + +P = ParamSpec("P") +R = TypeVar("R") +Ts = TypeVarTuple("Ts") +Us = TypeVarTuple("Us") + +class contextmanager(Generic[P, R]): + def __init__(self, func: Callable[P, R]) -> None: ... + + def __call__(self, *args: P.args, **kwargs: P.kwargs) -> "_contextmanager_cls[P, R]": ... + +class _contextmanager_cls(Generic[P, R]): + def __enter__(self) -> R: ... + def __exit__(self, *args: object) -> bool: ... + def invoke(self, *args: P.args, **kwargs: P.kwargs) -> R: ... + +@contextmanager +def print_args(*args: Unpack[Ts]) -> tuple[Unpack[Ts]]: ... + +with print_args(2, "x") as value: + reveal_type(value) # N: Revealed type is "tuple[builtins.int, builtins.str]" + +reveal_type(print_args(2, "x")) # N: Revealed type is "__main__._contextmanager_cls[[Literal[2]?, Literal['x']?], tuple[Literal[2]?, Literal['x']?]]" + +def forward(*args: Unpack[Us]) -> None: + manager = print_args(0, *args, "end") + reveal_type(manager.invoke(0, *args, "end")) # N: Revealed type is "tuple[builtins.int, Unpack[Us`-1], builtins.str]" +[builtins fixtures/tuple.pyi] + [case testInvalidParamSpecType] from typing import ParamSpec