Fix type inference for index expression with bounded TypeVar (#11434)

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.

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.

Author: ilyalabun

mypy/checkexpr.py

@@ -2964,6 +2964,9 @@ 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)
+              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(
                 '__getitem__', left_type, [e.index], [ARG_POS], e,

test-data/unit/check-typevar-values.test

@@ -631,3 +631,74 @@ 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]
+
+[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]
+
+[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]