Skip to content

Commit 713fa7b

Browse files
ankileankile
authored
Add type literal to engine argument in rolling functions (#54135)
1 parent b3648c4 commit 713fa7b

File tree

2 files changed

+26
-24
lines changed

2 files changed

+26
-24
lines changed

‎pandas/core/window/expanding.py

Lines changed: 9 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -5,6 +5,7 @@
55
TYPE_CHECKING,
66
Any,
77
Callable,
8+
Literal,
89
)
910

1011
from pandas.util._decorators import (
@@ -231,7 +232,7 @@ def apply(
231232
self,
232233
func: Callable[..., Any],
233234
raw: bool = False,
234-
engine: str | None = None,
235+
engine: Literal["cython", "numba"] | None = None,
235236
engine_kwargs: dict[str, bool] | None = None,
236237
args: tuple[Any, ...] | None = None,
237238
kwargs: dict[str, Any] | None = None,
@@ -275,7 +276,7 @@ def apply(
275276
def sum(
276277
self,
277278
numeric_only: bool = False,
278-
engine: str | None = None,
279+
engine: Literal["cython", "numba"] | None = None,
279280
engine_kwargs: dict[str, bool] | None = None,
280281
):
281282
return super().sum(
@@ -314,7 +315,7 @@ def sum(
314315
def max(
315316
self,
316317
numeric_only: bool = False,
317-
engine: str | None = None,
318+
engine: Literal["cython", "numba"] | None = None,
318319
engine_kwargs: dict[str, bool] | None = None,
319320
):
320321
return super().max(
@@ -353,7 +354,7 @@ def max(
353354
def min(
354355
self,
355356
numeric_only: bool = False,
356-
engine: str | None = None,
357+
engine: Literal["cython", "numba"] | None = None,
357358
engine_kwargs: dict[str, bool] | None = None,
358359
):
359360
return super().min(
@@ -392,7 +393,7 @@ def min(
392393
def mean(
393394
self,
394395
numeric_only: bool = False,
395-
engine: str | None = None,
396+
engine: Literal["cython", "numba"] | None = None,
396397
engine_kwargs: dict[str, bool] | None = None,
397398
):
398399
return super().mean(
@@ -431,7 +432,7 @@ def mean(
431432
def median(
432433
self,
433434
numeric_only: bool = False,
434-
engine: str | None = None,
435+
engine: Literal["cython", "numba"] | None = None,
435436
engine_kwargs: dict[str, bool] | None = None,
436437
):
437438
return super().median(
@@ -490,7 +491,7 @@ def std(
490491
self,
491492
ddof: int = 1,
492493
numeric_only: bool = False,
493-
engine: str | None = None,
494+
engine: Literal["cython", "numba"] | None = None,
494495
engine_kwargs: dict[str, bool] | None = None,
495496
):
496497
return super().std(
@@ -550,7 +551,7 @@ def var(
550551
self,
551552
ddof: int = 1,
552553
numeric_only: bool = False,
553-
engine: str | None = None,
554+
engine: Literal["cython", "numba"] | None = None,
554555
engine_kwargs: dict[str, bool] | None = None,
555556
):
556557
return super().var(

‎pandas/core/window/rolling.py

Lines changed: 17 additions & 16 deletions
Original file line numberDiff line numberDiff line change
@@ -13,6 +13,7 @@
1313
TYPE_CHECKING,
1414
Any,
1515
Callable,
16+
Literal,
1617
cast,
1718
)
1819

@@ -1458,7 +1459,7 @@ def apply(
14581459
self,
14591460
func: Callable[..., Any],
14601461
raw: bool = False,
1461-
engine: str | None = None,
1462+
engine: Literal["cython", "numba"] | None = None,
14621463
engine_kwargs: dict[str, bool] | None = None,
14631464
args: tuple[Any, ...] | None = None,
14641465
kwargs: dict[str, Any] | None = None,
@@ -1525,7 +1526,7 @@ def apply_func(values, begin, end, min_periods, raw=raw):
15251526
def sum(
15261527
self,
15271528
numeric_only: bool = False,
1528-
engine: str | None = None,
1529+
engine: Literal["cython", "numba"] | None = None,
15291530
engine_kwargs: dict[str, bool] | None = None,
15301531
):
15311532
if maybe_use_numba(engine):
@@ -1547,7 +1548,7 @@ def sum(
15471548
def max(
15481549
self,
15491550
numeric_only: bool = False,
1550-
engine: str | None = None,
1551+
engine: Literal["cython", "numba"] | None = None,
15511552
engine_kwargs: dict[str, bool] | None = None,
15521553
):
15531554
if maybe_use_numba(engine):
@@ -1569,7 +1570,7 @@ def max(
15691570
def min(
15701571
self,
15711572
numeric_only: bool = False,
1572-
engine: str | None = None,
1573+
engine: Literal["cython", "numba"] | None = None,
15731574
engine_kwargs: dict[str, bool] | None = None,
15741575
):
15751576
if maybe_use_numba(engine):
@@ -1591,7 +1592,7 @@ def min(
15911592
def mean(
15921593
self,
15931594
numeric_only: bool = False,
1594-
engine: str | None = None,
1595+
engine: Literal["cython", "numba"] | None = None,
15951596
engine_kwargs: dict[str, bool] | None = None,
15961597
):
15971598
if maybe_use_numba(engine):
@@ -1613,7 +1614,7 @@ def mean(
16131614
def median(
16141615
self,
16151616
numeric_only: bool = False,
1616-
engine: str | None = None,
1617+
engine: Literal["cython", "numba"] | None = None,
16171618
engine_kwargs: dict[str, bool] | None = None,
16181619
):
16191620
if maybe_use_numba(engine):
@@ -1635,7 +1636,7 @@ def std(
16351636
self,
16361637
ddof: int = 1,
16371638
numeric_only: bool = False,
1638-
engine: str | None = None,
1639+
engine: Literal["cython", "numba"] | None = None,
16391640
engine_kwargs: dict[str, bool] | None = None,
16401641
):
16411642
if maybe_use_numba(engine):
@@ -1659,7 +1660,7 @@ def var(
16591660
self,
16601661
ddof: int = 1,
16611662
numeric_only: bool = False,
1662-
engine: str | None = None,
1663+
engine: Literal["cython", "numba"] | None = None,
16631664
engine_kwargs: dict[str, bool] | None = None,
16641665
):
16651666
if maybe_use_numba(engine):
@@ -2021,7 +2022,7 @@ def apply(
20212022
self,
20222023
func: Callable[..., Any],
20232024
raw: bool = False,
2024-
engine: str | None = None,
2025+
engine: Literal["cython", "numba"] | None = None,
20252026
engine_kwargs: dict[str, bool] | None = None,
20262027
args: tuple[Any, ...] | None = None,
20272028
kwargs: dict[str, Any] | None = None,
@@ -2101,7 +2102,7 @@ def apply(
21012102
def sum(
21022103
self,
21032104
numeric_only: bool = False,
2104-
engine: str | None = None,
2105+
engine: Literal["cython", "numba"] | None = None,
21052106
engine_kwargs: dict[str, bool] | None = None,
21062107
):
21072108
return super().sum(
@@ -2141,7 +2142,7 @@ def max(
21412142
self,
21422143
numeric_only: bool = False,
21432144
*args,
2144-
engine: str | None = None,
2145+
engine: Literal["cython", "numba"] | None = None,
21452146
engine_kwargs: dict[str, bool] | None = None,
21462147
**kwargs,
21472148
):
@@ -2184,7 +2185,7 @@ def max(
21842185
def min(
21852186
self,
21862187
numeric_only: bool = False,
2187-
engine: str | None = None,
2188+
engine: Literal["cython", "numba"] | None = None,
21882189
engine_kwargs: dict[str, bool] | None = None,
21892190
):
21902191
return super().min(
@@ -2233,7 +2234,7 @@ def min(
22332234
def mean(
22342235
self,
22352236
numeric_only: bool = False,
2236-
engine: str | None = None,
2237+
engine: Literal["cython", "numba"] | None = None,
22372238
engine_kwargs: dict[str, bool] | None = None,
22382239
):
22392240
return super().mean(
@@ -2275,7 +2276,7 @@ def mean(
22752276
def median(
22762277
self,
22772278
numeric_only: bool = False,
2278-
engine: str | None = None,
2279+
engine: Literal["cython", "numba"] | None = None,
22792280
engine_kwargs: dict[str, bool] | None = None,
22802281
):
22812282
return super().median(
@@ -2333,7 +2334,7 @@ def std(
23332334
self,
23342335
ddof: int = 1,
23352336
numeric_only: bool = False,
2336-
engine: str | None = None,
2337+
engine: Literal["cython", "numba"] | None = None,
23372338
engine_kwargs: dict[str, bool] | None = None,
23382339
):
23392340
return super().std(
@@ -2392,7 +2393,7 @@ def var(
23922393
self,
23932394
ddof: int = 1,
23942395
numeric_only: bool = False,
2395-
engine: str | None = None,
2396+
engine: Literal["cython", "numba"] | None = None,
23962397
engine_kwargs: dict[str, bool] | None = None,
23972398
):
23982399
return super().var(

0 commit comments

Comments
 (0)