Plugin for tuple multiplication with literal int (#10361)

Author: hatal175

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

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"

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
 

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