diff --git a/mypy/plugins/default.py b/mypy/plugins/default.py index cb22b92c618f5..30f3b50a1b545 100644 --- a/mypy/plugins/default.py +++ b/mypy/plugins/default.py @@ -10,7 +10,7 @@ from mypy.plugins.common import try_getting_str_literals from mypy.types import ( FunctionLike, Type, Instance, AnyType, TypeOfAny, CallableType, NoneType, TypedDictType, - TypeVarType, TPDICT_FB_NAMES, get_proper_type, LiteralType + TypeVarType, TPDICT_FB_NAMES, get_proper_type, LiteralType, TupleType ) from mypy.subtypes import is_subtype from mypy.typeops import make_simplified_union @@ -64,6 +64,8 @@ def get_method_hook(self, fullname: str return int_pow_callback elif fullname == 'builtins.int.__neg__': return int_neg_callback + elif fullname in ('builtins.tuple.__mul__', 'builtins.tuple.__rmul__'): + return tuple_mul_callback elif fullname in set(n + '.setdefault' for n in TPDICT_FB_NAMES): return typed_dict_setdefault_callback elif fullname in set(n + '.pop' for n in TPDICT_FB_NAMES): @@ -471,3 +473,24 @@ def int_neg_callback(ctx: MethodContext) -> Type: if isinstance(value, int): return LiteralType(value=-value, fallback=fallback) return ctx.default_return_type + + +def tuple_mul_callback(ctx: MethodContext) -> Type: + """Infer a more precise return type for tuple.__mul__ and tuple.__rmul__. + + This is used to return a specific sized tuple if multiplied by Literal int + """ + if not isinstance(ctx.type, TupleType): + return ctx.default_return_type + + arg_type = ctx.arg_types[0][0] + if isinstance(arg_type, Instance) and arg_type.last_known_value is not None: + value = arg_type.last_known_value.value + if isinstance(value, int): + return ctx.type.copy_modified(items=ctx.type.items * value) + elif isinstance(ctx.type, LiteralType): + value = arg_type.value + if isinstance(value, int): + return ctx.type.copy_modified(items=ctx.type.items * value) + + return ctx.default_return_type diff --git a/test-data/unit/check-generics.test b/test-data/unit/check-generics.test index f048420b706ea..161b5d72e772b 100644 --- a/test-data/unit/check-generics.test +++ b/test-data/unit/check-generics.test @@ -1672,7 +1672,7 @@ def f(x: T) -> str: [case testTypeVarReversibleOperatorTuple] from typing import TypeVar, Tuple class A(Tuple[int, int]): - def __mul__(cls, other: Tuple[int, int]) -> str: return "" + def __mul__(cls, other: Tuple[int, int]) -> str: return "" # type: ignore # overriding default __mul__ T = TypeVar("T", bound=A) def f(x: T) -> str: return reveal_type(x * (1, 2) ) # N: Revealed type is "builtins.str" diff --git a/test-data/unit/check-tuples.test b/test-data/unit/check-tuples.test index ee71269726617..d5b35784f5e5f 100644 --- a/test-data/unit/check-tuples.test +++ b/test-data/unit/check-tuples.test @@ -1471,6 +1471,32 @@ x9, y9, x10, y10, z5 = *points2, 1, *points2 # E: Contiguous iterable with same () = 1 # E: "Literal[1]?" object is not iterable [builtins fixtures/tuple.pyi] +[case testMultiplyTupleByIntegerLiteral] +from typing import Tuple +t = ('',) * 2 +reveal_type(t) # N: Revealed type is "Tuple[builtins.str, builtins.str]" +t2 = ('',) * -1 +reveal_type(t2) # N: Revealed type is "Tuple[]" +t3 = ('', 1) * 2 +reveal_type(t3) # N: Revealed type is "Tuple[builtins.str, builtins.int, builtins.str, builtins.int]" +def f() -> Tuple[str, ...]: + return ('', ) +reveal_type(f() * 2) # N: Revealed type is "builtins.tuple[builtins.str*]" +[builtins fixtures/tuple.pyi] + +[case testMultiplyTupleByIntegerLiteralReverse] +from typing import Tuple +t = 2 * ('',) +reveal_type(t) # N: Revealed type is "Tuple[builtins.str, builtins.str]" +t2 = -1 * ('',) +reveal_type(t2) # N: Revealed type is "Tuple[]" +t3 = 2 * ('', 1) +reveal_type(t3) # N: Revealed type is "Tuple[builtins.str, builtins.int, builtins.str, builtins.int]" +def f() -> Tuple[str, ...]: + return ('', ) +reveal_type(2 * f()) # N: Revealed type is "builtins.tuple[builtins.str*]" +[builtins fixtures/tuple.pyi] + [case testSingleUndefinedTypeAndTuple] from typing import Tuple diff --git a/test-data/unit/fixtures/tuple.pyi b/test-data/unit/fixtures/tuple.pyi index a101595c6f30b..36c53de619bed 100644 --- a/test-data/unit/fixtures/tuple.pyi +++ b/test-data/unit/fixtures/tuple.pyi @@ -14,6 +14,7 @@ class tuple(Sequence[Tco], Generic[Tco]): def __iter__(self) -> Iterator[Tco]: pass def __contains__(self, item: object) -> bool: pass def __getitem__(self, x: int) -> Tco: pass + def __mul__(self, n: int) -> Tuple[Tco, ...]: pass def __rmul__(self, n: int) -> Tuple[Tco, ...]: pass def __add__(self, x: Tuple[Tco, ...]) -> Tuple[Tco, ...]: pass def count(self, obj: object) -> int: pass