From 59db820112c9fb701a1be677e049cdab9391c181 Mon Sep 17 00:00:00 2001 From: Ilya Labun Date: Tue, 2 Nov 2021 21:08:55 +0100 Subject: [PATCH 1/3] Fix type inference for index expression with bounded TypeVar Closes #8231 When type of index expression (e.g. `foo[bar]`) is analyzed and left expression (i.e. `foo`) has generic type (`TypeVar`) with upper bound, for some upper bound types this type inference yields wrong result. For example, if upper bound type is instance of `TypeDict`, mypy considers return type of such index expression as `object`: ``` from typing import TypedDict, TypeVar class Data(TypedDict): x: int T = TypeVar("T", bound=Data) def f(data: T) -> int: # line below leads to mypy error: # 'Unsupported operand types for + ("object" and "int")' return data["x"] + 1 ``` The root cause of this issue was in `visit_index_with_type` method from `checkexpr.py` which does type analysis for index expressions. For `TypeVar` left expression code flow goes via default branch which just returns return type of upper bound's `__getitem__`. For some types this return type inference logic operates on a fallback type. For example, for `TypedDict` fallback type is `typing._TypedDict` with `__getitem__` returning just `object`. To fix the issue we added special case to `visit_index_with_type` for `TypeVar` left expression which recursively calls `visit_index_with_type` with `TypeVar` upper bound parameter. This way we always handle upper bounds requiring special treatment correctly. --- mypy/checkexpr.py | 2 ++ test-data/unit/check-typevar-values.test | 15 +++++++++++++++ 2 files changed, 17 insertions(+) diff --git a/mypy/checkexpr.py b/mypy/checkexpr.py index b6dec2834449d..bc7299a1a046b 100644 --- a/mypy/checkexpr.py +++ b/mypy/checkexpr.py @@ -2960,6 +2960,8 @@ def visit_index_with_type(self, left_type: Type, e: IndexExpr, elif (isinstance(left_type, CallableType) and left_type.is_type_obj() and left_type.type_object().is_enum): return self.visit_enum_index_expr(left_type.type_object(), e.index, e) + elif isinstance(left_type, TypeVarType): + return self.visit_index_with_type(left_type.upper_bound, e, original_type) else: result, method_type = self.check_method_call_by_name( '__getitem__', left_type, [e.index], [ARG_POS], e, diff --git a/test-data/unit/check-typevar-values.test b/test-data/unit/check-typevar-values.test index 3f77996ec9596..a5b8f859b5945 100644 --- a/test-data/unit/check-typevar-values.test +++ b/test-data/unit/check-typevar-values.test @@ -631,3 +631,18 @@ def g(s: S) -> Callable[[S], None]: ... def f(x: S) -> None: h = g(x) h(x) + +[case testTypeVarWithTypedDictBoundInIndexExpression] +from typing import TypeVar +from typing_extensions import TypedDict + +class Data(TypedDict): + x: int + + +T = TypeVar("T", bound=Data) + + +def f(data: T) -> None: + reveal_type(data["x"]) # N: Revealed type is "builtins.int" +[builtins fixtures/tuple.pyi] From 8dc33159fb13883e4574731783f67f6ea049ab86 Mon Sep 17 00:00:00 2001 From: Ilya Labun Date: Wed, 3 Nov 2021 20:26:06 +0100 Subject: [PATCH 2/3] Add tests for TypeVar in Union and for TypeVar with values --- test-data/unit/check-typevar-values.test | 33 ++++++++++++++++++++++++ 1 file changed, 33 insertions(+) diff --git a/test-data/unit/check-typevar-values.test b/test-data/unit/check-typevar-values.test index a5b8f859b5945..145e708c50c6a 100644 --- a/test-data/unit/check-typevar-values.test +++ b/test-data/unit/check-typevar-values.test @@ -646,3 +646,36 @@ T = TypeVar("T", bound=Data) def f(data: T) -> None: reveal_type(data["x"]) # N: Revealed type is "builtins.int" [builtins fixtures/tuple.pyi] + +[case testTypeVarWithUnionTypedDictBoundInIndexExpression] +from typing import TypeVar, Union, Dict +from typing_extensions import TypedDict + +class Data(TypedDict): + x: int + + +T = TypeVar("T", bound=Union[Data, Dict[str, str]]) + + +def f(data: T) -> None: + reveal_type(data["x"]) # N: Revealed type is "Union[builtins.int, builtins.str*]" + +[builtins fixtures/tuple.pyi] +[builtins fixtures/dict.pyi] + +[case testTypeVarWithTypedDictValueInIndexExpression] +from typing import TypeVar, Union, Dict +from typing_extensions import TypedDict + +class Data(TypedDict): + x: int + + +T = TypeVar("T", Data, Dict[str, str]) + + +def f(data: T) -> None: + _: Union[str, int] = data["x"] +[builtins fixtures/tuple.pyi] +[builtins fixtures/dict.pyi] From d40b2b02dda040f410c678977cf03f3525f484be Mon Sep 17 00:00:00 2001 From: Ilya Labun Date: Sun, 7 Nov 2021 21:55:01 +0100 Subject: [PATCH 3/3] Corner case -- recursive TypeVar `T` Corner case -- recursive TypeVar `T` with upper bound having `__getitem__` method with `self` having type `T` and returning `T`. In this case when we call `visit_index_with_type` recursively, `TypeVar` types of upper bound `__getitem__` method are erased and `check_method_call_by_name` returns type of upper bound, not `T`. So we don't do recursive `visit_index_with_type` call when upper bound has `__getitem__` method and handle this case in `else` branch. which handles it as expected -- it return `T` as return type. --- mypy/checkexpr.py | 3 ++- test-data/unit/check-typevar-values.test | 23 +++++++++++++++++++++++ 2 files changed, 25 insertions(+), 1 deletion(-) diff --git a/mypy/checkexpr.py b/mypy/checkexpr.py index bc7299a1a046b..c550cf0e5f5a7 100644 --- a/mypy/checkexpr.py +++ b/mypy/checkexpr.py @@ -2960,7 +2960,8 @@ def visit_index_with_type(self, left_type: Type, e: IndexExpr, elif (isinstance(left_type, CallableType) and left_type.is_type_obj() and left_type.type_object().is_enum): return self.visit_enum_index_expr(left_type.type_object(), e.index, e) - elif isinstance(left_type, TypeVarType): + elif (isinstance(left_type, TypeVarType) + and not self.has_member(left_type.upper_bound, "__getitem__")): return self.visit_index_with_type(left_type.upper_bound, e, original_type) else: result, method_type = self.check_method_call_by_name( diff --git a/test-data/unit/check-typevar-values.test b/test-data/unit/check-typevar-values.test index 145e708c50c6a..2c25e9adc8bb8 100644 --- a/test-data/unit/check-typevar-values.test +++ b/test-data/unit/check-typevar-values.test @@ -679,3 +679,26 @@ def f(data: T) -> None: _: Union[str, int] = data["x"] [builtins fixtures/tuple.pyi] [builtins fixtures/dict.pyi] + +[case testSelfTypeVarIndexExpr] +from typing import TypeVar, Union, Type +from typing_extensions import TypedDict + +T = TypeVar("T", bound="Indexable") + +class Indexable: + def __init__(self, index: str) -> None: + self.index = index + + def __getitem__(self: T, index: str) -> T: + return self._new_instance(index) + + @classmethod + def _new_instance(cls: Type[T], index: str) -> T: + return cls("foo") + + def m(self: T) -> T: + return self["foo"] + +[builtins fixtures/tuple.pyi] +[builtins fixtures/classmethod.pyi]