From a625b806bdcc9bdb3f17cfd7ff5157b6896f7165 Mon Sep 17 00:00:00 2001 From: Tal Hayon Date: Fri, 23 Apr 2021 15:18:58 +0300 Subject: [PATCH 1/5] Add MethodCallback for Tuple and Literal multiplication --- mypy/plugins/default.py | 19 +++++++++++++++++++ test-data/unit/pythoneval.test | 18 ++++++++++++++++++ 2 files changed, 37 insertions(+) diff --git a/mypy/plugins/default.py b/mypy/plugins/default.py index 552a52c5c860d..23d8780f5e994 100644 --- a/mypy/plugins/default.py +++ b/mypy/plugins/default.py @@ -60,6 +60,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): @@ -463,3 +465,20 @@ 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 + """ + 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 \ No newline at end of file diff --git a/test-data/unit/pythoneval.test b/test-data/unit/pythoneval.test index 49f308af36100..141fbe45856f4 100644 --- a/test-data/unit/pythoneval.test +++ b/test-data/unit/pythoneval.test @@ -809,6 +809,24 @@ _program.py:3: note: Possible overload variants: _program.py:3: note: def __add__(self, Tuple[str, ...]) -> Tuple[str, ...] _program.py:3: note: def __add__(self, Tuple[Any, ...]) -> Tuple[Any, ...] +[case testMultiplyTupleByIntegerLiteral] +t = ('',) * 2 +reveal_type(t) +t2 = ('',) * -1 +reveal_type(t2) +[out] +_testMultiplyTupleByIntegerLiteral.py:2: note: Revealed type is "Tuple[builtins.str, builtins.str]" +_testMultiplyTupleByIntegerLiteral.py:4: note: Revealed type is "Tuple[]" + +[case testMultiplyTupleByIntegerLiteralReverse] +t = 2 * ('',) +reveal_type(t) +t2 = -1 * ('',) +reveal_type(t2) +[out] +_testMultiplyTupleByIntegerLiteralReverse.py:2: note: Revealed type is "Tuple[builtins.str, builtins.str]" +_testMultiplyTupleByIntegerLiteralReverse.py:4: note: Revealed type is "Tuple[]" + [case testDictWithKeywordArgs] from typing import Dict, Any, List d1 = dict(a=1, b=2) # type: Dict[str, int] From d7faadb0786e0306d9c5479234c0aa1147d900ed Mon Sep 17 00:00:00 2001 From: Tal Hayon Date: Fri, 23 Apr 2021 15:23:45 +0300 Subject: [PATCH 2/5] flake8 fixes --- mypy/plugins/default.py | 7 ++++--- 1 file changed, 4 insertions(+), 3 deletions(-) diff --git a/mypy/plugins/default.py b/mypy/plugins/default.py index 23d8780f5e994..e3ac4a7beaa11 100644 --- a/mypy/plugins/default.py +++ b/mypy/plugins/default.py @@ -466,6 +466,7 @@ def int_neg_callback(ctx: MethodContext) -> Type: 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__. @@ -475,10 +476,10 @@ def tuple_mul_callback(ctx: MethodContext) -> Type: 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) + 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.type.copy_modified(items=ctx.type.items * value) - return ctx.default_return_type \ No newline at end of file + return ctx.default_return_type From a8bccf0cda127d235eafa9b66db86e4b59802845 Mon Sep 17 00:00:00 2001 From: Tal Hayon Date: Fri, 23 Apr 2021 15:29:38 +0300 Subject: [PATCH 3/5] Add more test cases --- test-data/unit/pythoneval.test | 6 ++++++ 1 file changed, 6 insertions(+) diff --git a/test-data/unit/pythoneval.test b/test-data/unit/pythoneval.test index 141fbe45856f4..469104486e456 100644 --- a/test-data/unit/pythoneval.test +++ b/test-data/unit/pythoneval.test @@ -814,18 +814,24 @@ t = ('',) * 2 reveal_type(t) t2 = ('',) * -1 reveal_type(t2) +t3 = ('', 1) * 2 +reveal_type(t3) [out] _testMultiplyTupleByIntegerLiteral.py:2: note: Revealed type is "Tuple[builtins.str, builtins.str]" _testMultiplyTupleByIntegerLiteral.py:4: note: Revealed type is "Tuple[]" +_testMultiplyTupleByIntegerLiteral.py:6: note: Revealed type is "Tuple[builtins.str, builtins.int, builtins.str, builtins.int]" [case testMultiplyTupleByIntegerLiteralReverse] t = 2 * ('',) reveal_type(t) t2 = -1 * ('',) reveal_type(t2) +t3 = 2 * ('', 1) +reveal_type(t3) [out] _testMultiplyTupleByIntegerLiteralReverse.py:2: note: Revealed type is "Tuple[builtins.str, builtins.str]" _testMultiplyTupleByIntegerLiteralReverse.py:4: note: Revealed type is "Tuple[]" +_testMultiplyTupleByIntegerLiteralReverse.py:6: note: Revealed type is "Tuple[builtins.str, builtins.int, builtins.str, builtins.int]" [case testDictWithKeywordArgs] from typing import Dict, Any, List From 3dd35d320e099c3ec3724487cc5068aa2ed61b5b Mon Sep 17 00:00:00 2001 From: Tal Hayon Date: Fri, 23 Apr 2021 16:15:18 +0300 Subject: [PATCH 4/5] Fix handling variant size tuples --- mypy/plugins/default.py | 6 +++++- test-data/unit/pythoneval.test | 24 +++++++++++++++++------- 2 files changed, 22 insertions(+), 8 deletions(-) diff --git a/mypy/plugins/default.py b/mypy/plugins/default.py index e3ac4a7beaa11..9a373f07cb061 100644 --- a/mypy/plugins/default.py +++ b/mypy/plugins/default.py @@ -10,7 +10,8 @@ from mypy.plugins.common import try_getting_str_literals from mypy.types import ( Type, Instance, AnyType, TypeOfAny, CallableType, NoneType, TypedDictType, - TypeVarDef, TypeVarType, TPDICT_FB_NAMES, get_proper_type, LiteralType + TypeVarDef, TypeVarType, TPDICT_FB_NAMES, get_proper_type, LiteralType, + TupleType ) from mypy.subtypes import is_subtype from mypy.typeops import make_simplified_union @@ -472,6 +473,9 @@ def tuple_mul_callback(ctx: MethodContext) -> Type: 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 diff --git a/test-data/unit/pythoneval.test b/test-data/unit/pythoneval.test index 469104486e456..ea2a29d157223 100644 --- a/test-data/unit/pythoneval.test +++ b/test-data/unit/pythoneval.test @@ -810,28 +810,38 @@ _program.py:3: note: def __add__(self, Tuple[str, ...]) -> Tuple[str, ...] _program.py:3: note: def __add__(self, Tuple[Any, ...]) -> Tuple[Any, ...] [case testMultiplyTupleByIntegerLiteral] +from typing import Tuple t = ('',) * 2 reveal_type(t) t2 = ('',) * -1 reveal_type(t2) t3 = ('', 1) * 2 reveal_type(t3) +def f() -> Tuple[str, ...]: + return ('', ) +reveal_type(f() * 2) [out] -_testMultiplyTupleByIntegerLiteral.py:2: note: Revealed type is "Tuple[builtins.str, builtins.str]" -_testMultiplyTupleByIntegerLiteral.py:4: note: Revealed type is "Tuple[]" -_testMultiplyTupleByIntegerLiteral.py:6: note: Revealed type is "Tuple[builtins.str, builtins.int, builtins.str, builtins.int]" +_testMultiplyTupleByIntegerLiteral.py:3: note: Revealed type is "Tuple[builtins.str, builtins.str]" +_testMultiplyTupleByIntegerLiteral.py:5: note: Revealed type is "Tuple[]" +_testMultiplyTupleByIntegerLiteral.py:7: note: Revealed type is "Tuple[builtins.str, builtins.int, builtins.str, builtins.int]" +_testMultiplyTupleByIntegerLiteral.py:10: note: Revealed type is "builtins.tuple[builtins.str*]" [case testMultiplyTupleByIntegerLiteralReverse] +from typing import Tuple t = 2 * ('',) reveal_type(t) t2 = -1 * ('',) reveal_type(t2) t3 = 2 * ('', 1) reveal_type(t3) -[out] -_testMultiplyTupleByIntegerLiteralReverse.py:2: note: Revealed type is "Tuple[builtins.str, builtins.str]" -_testMultiplyTupleByIntegerLiteralReverse.py:4: note: Revealed type is "Tuple[]" -_testMultiplyTupleByIntegerLiteralReverse.py:6: note: Revealed type is "Tuple[builtins.str, builtins.int, builtins.str, builtins.int]" +def f() -> Tuple[str, ...]: + return ('', ) +reveal_type(2 * f()) +[out] +_testMultiplyTupleByIntegerLiteralReverse.py:3: note: Revealed type is "Tuple[builtins.str, builtins.str]" +_testMultiplyTupleByIntegerLiteralReverse.py:5: note: Revealed type is "Tuple[]" +_testMultiplyTupleByIntegerLiteralReverse.py:7: note: Revealed type is "Tuple[builtins.str, builtins.int, builtins.str, builtins.int]" +_testMultiplyTupleByIntegerLiteralReverse.py:10: note: Revealed type is "builtins.tuple[builtins.str*]" [case testDictWithKeywordArgs] from typing import Dict, Any, List From cd40202221cdbf156c39e950c1da6befc9ef6e7c Mon Sep 17 00:00:00 2001 From: Tal Hayon Date: Sat, 24 Apr 2021 05:09:11 +0300 Subject: [PATCH 5/5] Move tests to fixture tests --- test-data/unit/check-generics.test | 2 +- test-data/unit/check-tuples.test | 26 +++++++++++++++++++++++ test-data/unit/fixtures/tuple.pyi | 1 + test-data/unit/pythoneval.test | 34 ------------------------------ 4 files changed, 28 insertions(+), 35 deletions(-) diff --git a/test-data/unit/check-generics.test b/test-data/unit/check-generics.test index 651c2f3275024..8da2e16170b31 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 0018942df30cd..23981004bfdc8 100644 --- a/test-data/unit/check-tuples.test +++ b/test-data/unit/check-tuples.test @@ -1470,3 +1470,29 @@ x9, y9, x10, y10, z5 = *points2, 1, *points2 # E: Contiguous iterable with same [case testAssignEmptyBogus] () = 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] 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 diff --git a/test-data/unit/pythoneval.test b/test-data/unit/pythoneval.test index ea2a29d157223..49f308af36100 100644 --- a/test-data/unit/pythoneval.test +++ b/test-data/unit/pythoneval.test @@ -809,40 +809,6 @@ _program.py:3: note: Possible overload variants: _program.py:3: note: def __add__(self, Tuple[str, ...]) -> Tuple[str, ...] _program.py:3: note: def __add__(self, Tuple[Any, ...]) -> Tuple[Any, ...] -[case testMultiplyTupleByIntegerLiteral] -from typing import Tuple -t = ('',) * 2 -reveal_type(t) -t2 = ('',) * -1 -reveal_type(t2) -t3 = ('', 1) * 2 -reveal_type(t3) -def f() -> Tuple[str, ...]: - return ('', ) -reveal_type(f() * 2) -[out] -_testMultiplyTupleByIntegerLiteral.py:3: note: Revealed type is "Tuple[builtins.str, builtins.str]" -_testMultiplyTupleByIntegerLiteral.py:5: note: Revealed type is "Tuple[]" -_testMultiplyTupleByIntegerLiteral.py:7: note: Revealed type is "Tuple[builtins.str, builtins.int, builtins.str, builtins.int]" -_testMultiplyTupleByIntegerLiteral.py:10: note: Revealed type is "builtins.tuple[builtins.str*]" - -[case testMultiplyTupleByIntegerLiteralReverse] -from typing import Tuple -t = 2 * ('',) -reveal_type(t) -t2 = -1 * ('',) -reveal_type(t2) -t3 = 2 * ('', 1) -reveal_type(t3) -def f() -> Tuple[str, ...]: - return ('', ) -reveal_type(2 * f()) -[out] -_testMultiplyTupleByIntegerLiteralReverse.py:3: note: Revealed type is "Tuple[builtins.str, builtins.str]" -_testMultiplyTupleByIntegerLiteralReverse.py:5: note: Revealed type is "Tuple[]" -_testMultiplyTupleByIntegerLiteralReverse.py:7: note: Revealed type is "Tuple[builtins.str, builtins.int, builtins.str, builtins.int]" -_testMultiplyTupleByIntegerLiteralReverse.py:10: note: Revealed type is "builtins.tuple[builtins.str*]" - [case testDictWithKeywordArgs] from typing import Dict, Any, List d1 = dict(a=1, b=2) # type: Dict[str, int]