python / mypy

Optional static typing for Python
https://www.mypy-lang.org/
Other
18.53k stars 2.83k forks source link

Infer generic type arguments for slice expressions #18160

Closed brianschubert closed 3 days ago

brianschubert commented 4 days ago

Fixes #18149

Slices were made generic in https://github.com/python/typeshed/pull/11637. Currently, all slice expressions are inferred to have type slice[Any, Any, Any]. This PR fills in the generic type arguments more appropriately using start/stop/stride expression types.

Given

class Foo:
    def __getitem__[T](self, item: T) -> T: return item

x = Foo()
reveal_type(x[1:])

Before:

main.py:5: note: Revealed type is "builtins.slice[Any, Any, Any]"

After:

main.py:5: note: Revealed type is "builtins.slice[builtins.int, None, None]"
brianschubert commented 4 days ago

As an alternative approach, I toyed with using the slice.__new__ overloads from typeshed to infer the type arguments. However, this had a big impact on the tests (since all tests involving slice expressions now needed a fixture defining slice), and I'm not sure if there are any benefits from doing it that way.

```diff --- a/mypy/checkexpr.py +++ b/mypy/checkexpr.py @@ -5616,7 +5616,12 @@ class ExpressionChecker(ExpressionVisitor[Type]): if index: t = self.accept(index) self.chk.check_subtype(t, expected, index, message_registry.INVALID_SLICE_INDEX) - return self.named_type("builtins.slice") + args = [ + TempNode(NoneType()) if arg is None else arg + for arg in [e.begin_index, e.end_index, e.stride] + ] + slice_type = TypeType(self.named_type("builtins.slice")) + return self.check_call(slice_type, args, [nodes.ARG_POS] * 3, e, [None] * 3)[0] ```
github-actions[bot] commented 4 days ago

Diff from mypy_primer, showing the effect of this PR on open source code:

pandas-stubs (https://github.com/pandas-dev/pandas-stubs)
+ tests/test_frame.py:2246: error: Expression is of type "tuple[Index[int], slice[None, None, None]]", not "tuple[Index[int], slice[Any, Any, Any]]"  [assert-type]
- tests/test_frame.py:226: error: Expression is of type "Any", not "DataFrame"  [assert-type]
- tests/test_frame.py:227: error: Expression is of type "Any", not "DataFrame"  [assert-type]
brianschubert commented 4 days ago

Sure, done!

(FWIW, that's how I originally wrote it :wink:. I thought to try separating this logic from that loop since I figured it may be removed soon, now that generic slices let us do better checking on slice indices than the hardcoded SupportsIndex/None checks in that loop)

github-actions[bot] commented 4 days ago

Diff from mypy_primer, showing the effect of this PR on open source code:

pandas-stubs (https://github.com/pandas-dev/pandas-stubs)
+ tests/test_frame.py:2246: error: Expression is of type "tuple[Index[int], slice[None, None, None]]", not "tuple[Index[int], slice[Any, Any, Any]]"  [assert-type]
- tests/test_frame.py:226: error: Expression is of type "Any", not "DataFrame"  [assert-type]
- tests/test_frame.py:227: error: Expression is of type "Any", not "DataFrame"  [assert-type]