Repository: python/mypy
Pull request: 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.

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.

Problem statement:
TypeDict in bound TypeVar assumes parameter values are of type "object"
# Bug Report?

I may be misunderstanding how `TypeVar` works, but it seems to me that the following usage with `TypeDict` should be valid, however I get the error indicated in the comment:

```python
from typing import TypedDict, TypeVar


class Data(TypedDict):
    x: int


T = TypeVar("T", bound=Data)


def f(data: T) -> int:
    return data["x"] + 1  # E: Unsupported operand types for + ("object" and "int")
    

assert f(Data(x=1)) == 2
```

## Report Details

- have you tried it on `master`: no
- what version of MyPy: 0.761
- what flags are you using: none

Hints from issue discussion:
This does look like a bug, yeah. I think the issue is that when doing indexing on a type variable with a typeddict bound, checks it just as a call to `__getitem__` and ends up operating on the typeddict's "fallback" type, which is `typing._TypedDict`, which inherits from `Mapping[str, object]`.

I think it should be pretty easy to make `visit_index_with_type` in `checkexpr.py` try recursively indexing into its bound, if you want to take a shot at fixing it.
I'll take a shot at it.
I believe I'm having this issue with non typed dict too (if that's intended).  Some code to reproduce:

```python
from typing import Dict, List

someList = ["a", "b", "c", "c", "d"]
d: Dict[str, int] = {}

for x in someList:
    if v := d.get(x):
        v += 1 # Unsupported operand types for + ("None" and "int")mypy(error)
        d[x] = v
    else:
        d[x] = 1
```

Using VSCode, python==3.8.1, mypy==0.761
@np-hg that's a different issue, probably related to type narrowing.
Should I move it to a new issue? I didn't want to open one without being sure its not a duplicate.
@np-hg Your issue seems to be fixed on GitHub master.

Commit message: 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.

Patch to review:
diff --git a/mypy/checkexpr.py b/mypy/checkexpr.py
index b6dec2834449..bc7299a1a046 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 3f77996ec959..a5b8f859b594 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]


Code review comments:
diff_hunk: @@ -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)
path: mypy/checkexpr.py
text: @user1:
What if this type var has `.values` instead of `.upper_bound`? Will it still work?

@author:
surprisingly it works. I'm not sure why, but I suspect it's related to `TypeVar`s expansion in `checker.py::check_func_def`. Added test for this case as well