Skip to content

Commit be17dc0

Browse files
builtins.slice: more precise __new__ overloads and defaults for StopT and StepT. (#13008)
1 parent cadaaad commit be17dc0

File tree

2 files changed

+281
-14
lines changed

2 files changed

+281
-14
lines changed
Lines changed: 251 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,251 @@
1+
"""
2+
Assuming X, Y and Z are types other than None, the following rules apply to the slice type:
3+
4+
- The type hint `slice` should be compatible with all slices, including:
5+
- `slice(None)`, `slice(None, None)` and `slice(None, None, None)`. (⟿ `slice[?, ?, ?]`)
6+
- The type hint `slice[T]` should be compatible with:
7+
- `slice(None)`, `slice(None, None)` and `slice(None, None, None)` (⟿ `slice[?, ?, ?]`)
8+
- `slice(t)`, `slice(None, t)` and `slice(None, t, None)`. (⟿ `slice[?, T, ?]`)
9+
- `slice(t, None)` and `slice(t, None, None)`. (⟿ `slice[T, ?, ?]`)
10+
- `slice(t, t)` and `slice(t, t, None)`. (⟿ `slice[T, T, ?]`)
11+
- The type hint `slice[X, Y]` should be compatible with:
12+
- `slice(None)`, `slice(None, None)` and `slice(None, None, None)` (⟿ `slice[?, ?, ?]`)
13+
- `slice(y)`, `slice(None, y)` and `slice(None, y, None)`. (⟿ `slice[?, Y, ?]`)
14+
- `slice(x, None)` and `slice(x, None, None)` (⟿ `slice[X, ?, ?]`)
15+
- `slice(x, y)` and `slice(x, y, None)`. (⟿ `slice[X, Y, ?]`)
16+
- The type hint `slice[X, Y, Z]` should be compatible with:
17+
- `slice(None)`, `slice(None, None)` and `slice(None, None, None)`. (⟿ `slice[?, ?, ?]`)
18+
- `slice(y)`, `slice(None, y)` and `slice(None, y, None)`. (⟿ `slice[?, Y, ?]`)
19+
- `slice(x, None)` and `slice(x, None, None)` (⟿ `slice[X, ?, ?]`)
20+
- `slice(x, y)` and `slice(x, y, None)`. (⟿ `slice[X, Y, ?]`)
21+
- `slice(None, None, z)` (⟿ `slice[?, ?, Z]`)
22+
- `slice(None, y, z)` (⟿ `slice[?, Y, Z]`)
23+
- `slice(x, None, z)` (⟿ `slice[X, ?, Z]`)
24+
- `slice(x, y, z)` (⟿ `slice[X, Y, Z]`)
25+
26+
Consistency criterion: Assuming now X, Y, Z can potentially be None, the following rules apply:
27+
28+
- `slice(x)` must be compatible with `slice[None, X, None]`, even if X is None.
29+
- `slice(x, y)` must be compatible with `slice[X,Y,None]`, even if X is None or Y is None.
30+
- `slice(x, y, z)` must be compatible with `slice[X, Y, Z]`, even if X, Y, or Z are `None`.
31+
"""
32+
33+
from __future__ import annotations
34+
35+
from datetime import date, datetime as DT, timedelta as TD
36+
from typing import Any, SupportsIndex, cast
37+
from typing_extensions import assert_type
38+
39+
# region Tests for slice constructor overloads -----------------------------------------
40+
assert_type(slice(None), "slice[Any, Any, Any]")
41+
assert_type(slice(1234), "slice[Any, int, Any]")
42+
43+
assert_type(slice(None, None), "slice[Any, Any, Any]")
44+
assert_type(slice(None, 5678), "slice[Any, int, Any]")
45+
assert_type(slice(1234, None), "slice[int, Any, Any]")
46+
assert_type(slice(1234, 5678), "slice[int, int, Any]")
47+
48+
assert_type(slice(None, None, None), "slice[Any, Any, Any]")
49+
assert_type(slice(None, 5678, None), "slice[Any, int, Any]")
50+
assert_type(slice(1234, None, None), "slice[int, Any, Any]")
51+
assert_type(slice(1234, 5678, None), "slice[int, int, Any]")
52+
assert_type(slice(1234, 5678, 9012), "slice[int, int, int]")
53+
# endregion Tests for slice constructor overloads --------------------------------------
54+
55+
# region Test parameter defaults for slice constructor ---------------------------------
56+
# Note: need to cast, because pyright specializes regardless of type annotations
57+
slc1 = cast("slice[SupportsIndex | None]", slice(1))
58+
slc2 = cast("slice[int | None, int | None]", slice(1, 2))
59+
fake_key_val = cast("slice[str, int]", slice("1", 2))
60+
assert_type(slc1, "slice[SupportsIndex | None, SupportsIndex | None, SupportsIndex | None]")
61+
assert_type(slc2, "slice[int | None, int | None, int | None]")
62+
assert_type(fake_key_val, "slice[str, int, str | int]")
63+
# endregion Test parameter defaults for slice constructor ------------------------------
64+
65+
# region Tests for slice properties ----------------------------------------------------
66+
# Note: if an argument is not None, we should get precisely the same type back
67+
assert_type(slice(1234).stop, int)
68+
69+
assert_type(slice(1234, None).start, int)
70+
assert_type(slice(None, 5678).stop, int)
71+
72+
assert_type(slice(1234, None, None).start, int)
73+
assert_type(slice(None, 5678, None).stop, int)
74+
assert_type(slice(None, None, 9012).step, int)
75+
# endregion Tests for slice properties -------------------------------------------------
76+
77+
78+
# region Test for slice assignments ----------------------------------------------------
79+
# exhaustively test all possible assignments: miss (X), None (N), int (I), and str (S)
80+
rXNX: slice = slice(None)
81+
rXIX: slice = slice(1234)
82+
rXSX: slice = slice("70")
83+
84+
rNNX: slice = slice(None, None)
85+
rINX: slice = slice(1234, None)
86+
rSNX: slice = slice("70", None)
87+
88+
rNIX: slice = slice(None, 5678)
89+
rIIX: slice = slice(1234, 5678)
90+
rSIX: slice = slice("70", 9012)
91+
92+
rNSX: slice = slice(None, "71")
93+
rISX: slice = slice(1234, "71")
94+
rSSX: slice = slice("70", "71")
95+
96+
rNNN: slice = slice(None, None, None)
97+
rINN: slice = slice(1234, None, None)
98+
rSNN: slice = slice("70", None, None)
99+
rNIN: slice = slice(None, 5678, None)
100+
rIIN: slice = slice(1234, 5678, None)
101+
rSIN: slice = slice("70", 5678, None)
102+
rNSN: slice = slice(None, "71", None)
103+
rISN: slice = slice(1234, "71", None)
104+
rSSN: slice = slice("70", "71", None)
105+
106+
rNNI: slice = slice(None, None, 9012)
107+
rINI: slice = slice(1234, None, 9012)
108+
rSNI: slice = slice("70", None, 9012)
109+
rNII: slice = slice(None, 5678, 9012)
110+
rIII: slice = slice(1234, 5678, 9012)
111+
rSII: slice = slice("70", 5678, 9012)
112+
rNSI: slice = slice(None, "71", 9012)
113+
rISI: slice = slice(1234, "71", 9012)
114+
rSSI: slice = slice("70", "71", 9012)
115+
116+
rNNS: slice = slice(None, None, "1d")
117+
rINS: slice = slice(1234, None, "1d")
118+
rSNS: slice = slice("70", None, "1d")
119+
rNIS: slice = slice(None, 5678, "1d")
120+
rIIS: slice = slice(1234, 5678, "1d")
121+
rSIS: slice = slice("70", 5678, "1d")
122+
rNSS: slice = slice(None, "71", "1d")
123+
rISS: slice = slice(1234, "71", "1d")
124+
rSSS: slice = slice("70", "71", "1d")
125+
# endregion Test for slice assignments -------------------------------------------------
126+
127+
128+
# region Tests for slice[T] assignments ------------------------------------------------
129+
sXNX: "slice[int]" = slice(None)
130+
sXIX: "slice[int]" = slice(1234)
131+
132+
sNNX: "slice[int]" = slice(None, None)
133+
sNIX: "slice[int]" = slice(None, 5678)
134+
sINX: "slice[int]" = slice(1234, None)
135+
sIIX: "slice[int]" = slice(1234, 5678)
136+
137+
sNNN: "slice[int]" = slice(None, None, None)
138+
sNIN: "slice[int]" = slice(None, 5678, None)
139+
sNNS: "slice[int]" = slice(None, None, 9012)
140+
sINN: "slice[int]" = slice(1234, None, None)
141+
sINS: "slice[int]" = slice(1234, None, 9012)
142+
sIIN: "slice[int]" = slice(1234, 5678, None)
143+
sIIS: "slice[int]" = slice(1234, 5678, 9012)
144+
# endregion Tests for slice[T] assignments ---------------------------------------------
145+
146+
147+
# region Tests for slice[X, Y] assignments ---------------------------------------------
148+
# Note: start=int is illegal and hence we add an explicit "type: ignore" comment.
149+
tXNX: "slice[None, int]" = slice(None) # since slice(None) is slice[Any, Any, Any]
150+
tXIX: "slice[None, int]" = slice(1234)
151+
152+
tNNX: "slice[None, int]" = slice(None, None)
153+
tNIX: "slice[None, int]" = slice(None, 5678)
154+
tINX: "slice[None, int]" = slice(1234, None) # type: ignore
155+
tIIX: "slice[None, int]" = slice(1234, 5678) # type: ignore
156+
157+
tNNN: "slice[None, int]" = slice(None, None, None)
158+
tNIN: "slice[None, int]" = slice(None, 5678, None)
159+
tINN: "slice[None, int]" = slice(1234, None, None) # type: ignore
160+
tIIN: "slice[None, int]" = slice(1234, 5678, None) # type: ignore
161+
tNNS: "slice[None, int]" = slice(None, None, 9012)
162+
tINS: "slice[None, int]" = slice(None, 5678, 9012)
163+
tNIS: "slice[None, int]" = slice(1234, None, 9012) # type: ignore
164+
tIIS: "slice[None, int]" = slice(1234, 5678, 9012) # type: ignore
165+
# endregion Tests for slice[X, Y] assignments ------------------------------------------
166+
167+
168+
# region Tests for slice[X, Y, Z] assignments ------------------------------------------
169+
uXNX: "slice[int, int, int]" = slice(None)
170+
uXIX: "slice[int, int, int]" = slice(1234)
171+
172+
uNNX: "slice[int, int, int]" = slice(None, None)
173+
uNIX: "slice[int, int, int]" = slice(None, 5678)
174+
uINX: "slice[int, int, int]" = slice(1234, None)
175+
uIIX: "slice[int, int, int]" = slice(1234, 5678)
176+
177+
uNNN: "slice[int, int, int]" = slice(None, None, None)
178+
uNNI: "slice[int, int, int]" = slice(None, None, 9012)
179+
uNIN: "slice[int, int, int]" = slice(None, 5678, None)
180+
uNII: "slice[int, int, int]" = slice(None, 5678, 9012)
181+
uINN: "slice[int, int, int]" = slice(1234, None, None)
182+
uINI: "slice[int, int, int]" = slice(1234, None, 9012)
183+
uIIN: "slice[int, int, int]" = slice(1234, 5678, None)
184+
uIII: "slice[int, int, int]" = slice(1234, 5678, 9012)
185+
# endregion Tests for slice[X, Y, Z] assignments ---------------------------------------
186+
187+
188+
# region Test for slice consistency criterion ------------------------------------------
189+
year = date(2021, 1, 1)
190+
vXNX: "slice[None, None, None]" = slice(None)
191+
vXIX: "slice[None, date, None]" = slice(year)
192+
193+
vNNX: "slice[None, None, None]" = slice(None, None)
194+
vNIX: "slice[None, date, None]" = slice(None, year)
195+
vINX: "slice[date, None, None]" = slice(year, None)
196+
vIIX: "slice[date, date, None]" = slice(year, year)
197+
198+
vNNN: "slice[None, None, None]" = slice(None, None, None)
199+
vNIN: "slice[None, date, None]" = slice(None, year, None)
200+
vINN: "slice[date, None, None]" = slice(year, None, None)
201+
vIIN: "slice[date, date, None]" = slice(year, year, None)
202+
vNNI: "slice[None, None, str]" = slice(None, None, "1d")
203+
vNII: "slice[None, date, str]" = slice(None, year, "1d")
204+
vINI: "slice[date, None, str]" = slice(year, None, "1d")
205+
vIII: "slice[date, date, str]" = slice(year, year, "1d")
206+
# endregion Test for slice consistency criterion ---------------------------------------
207+
208+
209+
# region Integration tests for slices with datetimes -----------------------------------
210+
class TimeSeries: # similar to pandas.Series with datetime index
211+
def __getitem__(self, key: "slice[DT | str | None, DT | str | None]") -> Any:
212+
"""Subsample the time series at the given dates."""
213+
...
214+
215+
216+
class TimeSeriesInterpolator: # similar to pandas.Series with datetime index
217+
def __getitem__(self, key: "slice[DT, DT, TD | None]") -> Any:
218+
"""Subsample the time series at the given dates."""
219+
...
220+
221+
222+
# tests slices as an argument
223+
start = DT(1970, 1, 1)
224+
stop = DT(1971, 1, 10)
225+
step = TD(days=1)
226+
# see: https://pandas.pydata.org/docs/user_guide/timeseries.html#partial-string-indexing
227+
# FIXME: https://github.com/python/mypy/issues/2410 (use literal slices)
228+
series = TimeSeries()
229+
_ = series[slice(None, "1970-01-10")]
230+
_ = series[slice("1970-01-01", None)]
231+
_ = series[slice("1970-01-01", "1971-01-10")]
232+
_ = series[slice(None, stop)]
233+
_ = series[slice(start, None)]
234+
_ = series[slice(start, stop)]
235+
_ = series[slice(None)]
236+
237+
model = TimeSeriesInterpolator()
238+
_ = model[slice(start, stop)]
239+
_ = model[slice(start, stop, step)]
240+
_ = model[slice(start, stop, None)]
241+
242+
243+
# test slices as a return type
244+
def foo(flag: bool, value: DT) -> "slice[DT, None] | slice[None, DT]":
245+
if flag:
246+
return slice(value, None) # slice[DT, DT|Any, Any] incompatible
247+
else:
248+
return slice(None, value) # slice[DT|Any, DT, Any] incompatible
249+
250+
251+
# endregion Integration tests for slices with datetimes --------------------------------

stdlib/builtins.pyi

Lines changed: 30 additions & 14 deletions
Original file line numberDiff line numberDiff line change
@@ -10,7 +10,6 @@ from _typeshed import (
1010
ConvertibleToFloat,
1111
ConvertibleToInt,
1212
FileDescriptorOrPath,
13-
MaybeNone,
1413
OpenBinaryMode,
1514
OpenBinaryModeReading,
1615
OpenBinaryModeUpdating,
@@ -95,9 +94,14 @@ _SupportsAnextT = TypeVar("_SupportsAnextT", bound=SupportsAnext[Any], covariant
9594
_AwaitableT = TypeVar("_AwaitableT", bound=Awaitable[Any])
9695
_AwaitableT_co = TypeVar("_AwaitableT_co", bound=Awaitable[Any], covariant=True)
9796
_P = ParamSpec("_P")
98-
_StartT = TypeVar("_StartT", covariant=True, default=Any)
99-
_StopT = TypeVar("_StopT", covariant=True, default=Any)
100-
_StepT = TypeVar("_StepT", covariant=True, default=Any)
97+
98+
# Type variables for slice
99+
_StartT_co = TypeVar("_StartT_co", covariant=True, default=Any) # slice -> slice[Any, Any, Any]
100+
_StopT_co = TypeVar("_StopT_co", covariant=True, default=_StartT_co) # slice[A] -> slice[A, A, A]
101+
# NOTE: step could differ from start and stop, (e.g. datetime/timedelta)l
102+
# the default (start|stop) is chosen to cater to the most common case of int/index slices.
103+
# FIXME: https://github.com/python/typing/issues/213 (replace step=start|stop with step=start&stop)
104+
_StepT_co = TypeVar("_StepT_co", covariant=True, default=_StartT_co | _StopT_co) # slice[A,B] -> slice[A, B, A|B]
101105

102106
class object:
103107
__doc__: str | None
@@ -940,23 +944,35 @@ class bool(int):
940944
def __invert__(self) -> int: ...
941945

942946
@final
943-
class slice(Generic[_StartT, _StopT, _StepT]):
947+
class slice(Generic[_StartT_co, _StopT_co, _StepT_co]):
944948
@property
945-
def start(self) -> _StartT: ...
949+
def start(self) -> _StartT_co: ...
946950
@property
947-
def step(self) -> _StepT: ...
951+
def step(self) -> _StepT_co: ...
948952
@property
949-
def stop(self) -> _StopT: ...
950-
@overload
951-
def __new__(cls, stop: int | None, /) -> slice[int | MaybeNone, int | MaybeNone, int | MaybeNone]: ...
953+
def stop(self) -> _StopT_co: ...
954+
# Note: __new__ overloads map `None` to `Any`, since users expect slice(x, None)
955+
# to be compatible with slice(None, x).
956+
# generic slice --------------------------------------------------------------------
952957
@overload
953-
def __new__(
954-
cls, start: int | None, stop: int | None, step: int | None = None, /
955-
) -> slice[int | MaybeNone, int | MaybeNone, int | MaybeNone]: ...
958+
def __new__(cls, start: None, stop: None = None, step: None = None, /) -> slice[Any, Any, Any]: ...
959+
# unary overloads ------------------------------------------------------------------
956960
@overload
957961
def __new__(cls, stop: _T2, /) -> slice[Any, _T2, Any]: ...
962+
# binary overloads -----------------------------------------------------------------
963+
@overload
964+
def __new__(cls, start: _T1, stop: None, step: None = None, /) -> slice[_T1, Any, Any]: ...
965+
@overload
966+
def __new__(cls, start: None, stop: _T2, step: None = None, /) -> slice[Any, _T2, Any]: ...
967+
@overload
968+
def __new__(cls, start: _T1, stop: _T2, step: None = None, /) -> slice[_T1, _T2, Any]: ...
969+
# ternary overloads ----------------------------------------------------------------
970+
@overload
971+
def __new__(cls, start: None, stop: None, step: _T3, /) -> slice[Any, Any, _T3]: ...
972+
@overload
973+
def __new__(cls, start: _T1, stop: None, step: _T3, /) -> slice[_T1, Any, _T3]: ...
958974
@overload
959-
def __new__(cls, start: _T1, stop: _T2, /) -> slice[_T1, _T2, Any]: ...
975+
def __new__(cls, start: None, stop: _T2, step: _T3, /) -> slice[Any, _T2, _T3]: ...
960976
@overload
961977
def __new__(cls, start: _T1, stop: _T2, step: _T3, /) -> slice[_T1, _T2, _T3]: ...
962978
def __eq__(self, value: object, /) -> bool: ...

0 commit comments

Comments
 (0)