Skip to content

Commit d4d399e

Browse files
GH1055 Add pandas.api.typing to pandas-stubs (#1058)
* GH1055 Add pandas.api.typing to pandas-stubs * GH1055 Fix lint * GH1055 Fix assert_type * GH1055 Fix tests * Revert to generic * GH1055 Fix tests * GH1055 new test format * GH1055 PR Feedback * GH1055 TypeAlias Feedback
1 parent ddccc56 commit d4d399e

File tree

2 files changed

+247
-0
lines changed

2 files changed

+247
-0
lines changed

pandas-stubs/api/typing/__init__.pyi

Lines changed: 29 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,29 @@
1+
from pandas.core.groupby import (
2+
DataFrameGroupBy as DataFrameGroupBy,
3+
SeriesGroupBy as SeriesGroupBy,
4+
)
5+
from pandas.core.resample import (
6+
DatetimeIndexResamplerGroupby as DatetimeIndexResamplerGroupby,
7+
PeriodIndexResamplerGroupby as PeriodIndexResamplerGroupby,
8+
Resampler as Resampler,
9+
TimedeltaIndexResamplerGroupby as TimedeltaIndexResamplerGroupby,
10+
TimeGrouper as TimeGrouper,
11+
)
12+
from pandas.core.window import (
13+
Expanding as Expanding,
14+
ExpandingGroupby as ExpandingGroupby,
15+
ExponentialMovingWindow as ExponentialMovingWindow,
16+
ExponentialMovingWindowGroupby as ExponentialMovingWindowGroupby,
17+
Rolling as Rolling,
18+
RollingGroupby as RollingGroupby,
19+
Window as Window,
20+
)
21+
22+
from pandas._libs import NaTType as NaTType
23+
from pandas._libs.missing import NAType as NAType
24+
25+
from pandas.io.json._json import JsonReader as JsonReader
26+
27+
# SASReader is not defined so commenting it out for now
28+
# from pandas.io.sas.sasreader import SASReader as SASReader
29+
from pandas.io.stata import StataReader as StataReader

tests/test_api_typing.py

Lines changed: 218 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,218 @@
1+
"""Test module for classes in pandas.api.typing."""
2+
3+
import numpy as np
4+
import pandas as pd
5+
from pandas._testing import ensure_clean
6+
from pandas.api.typing import (
7+
DataFrameGroupBy,
8+
DatetimeIndexResamplerGroupby,
9+
Expanding,
10+
ExpandingGroupby,
11+
ExponentialMovingWindow,
12+
ExponentialMovingWindowGroupby,
13+
JsonReader,
14+
NaTType,
15+
NAType,
16+
PeriodIndexResamplerGroupby,
17+
Resampler,
18+
Rolling,
19+
RollingGroupby,
20+
SeriesGroupBy,
21+
StataReader,
22+
TimedeltaIndexResamplerGroupby,
23+
TimeGrouper,
24+
Window,
25+
)
26+
import pytest
27+
from typing_extensions import (
28+
TypeAlias,
29+
assert_type,
30+
)
31+
32+
from tests import check
33+
34+
from pandas.io.json._json import read_json
35+
36+
ResamplerGroupBy: TypeAlias = (
37+
DatetimeIndexResamplerGroupby
38+
| PeriodIndexResamplerGroupby
39+
| TimedeltaIndexResamplerGroupby
40+
)
41+
42+
43+
def test_dataframegroupby():
44+
df = pd.DataFrame({"a": [1, 2, 3]})
45+
group = df.groupby("a")
46+
47+
def f1(gb: DataFrameGroupBy):
48+
check(gb, DataFrameGroupBy)
49+
50+
f1(group)
51+
52+
53+
def test_seriesgroupby():
54+
sr = pd.Series([1, 2, 3], index=pd.Index(["a", "b", "a"]))
55+
56+
def f1(gb: SeriesGroupBy):
57+
check(gb, SeriesGroupBy)
58+
59+
f1(sr.groupby(level=0))
60+
61+
62+
def tests_datetimeindexersamplergroupby() -> None:
63+
idx = pd.date_range("1999-1-1", periods=365, freq="D")
64+
df = pd.DataFrame(
65+
np.random.standard_normal((365, 2)), index=idx, columns=["col1", "col2"]
66+
)
67+
gb_df = df.groupby("col2")
68+
69+
def f1(gb: ResamplerGroupBy):
70+
check(gb, DatetimeIndexResamplerGroupby)
71+
72+
f1(gb_df.resample("ME"))
73+
74+
75+
def test_timedeltaindexresamplergroupby() -> None:
76+
idx = pd.TimedeltaIndex(["0 days", "1 days", "2 days", "3 days", "4 days"])
77+
df = pd.DataFrame(
78+
np.random.standard_normal((5, 2)), index=idx, columns=["col1", "col2"]
79+
)
80+
gb_df = df.groupby("col2")
81+
82+
def f1(gb: ResamplerGroupBy):
83+
check(gb, TimedeltaIndexResamplerGroupby)
84+
85+
f1(gb_df.resample("1D"))
86+
87+
88+
@pytest.mark.skip("Resampling with a PeriodIndex is deprecated.")
89+
def test_periodindexresamplergroupby() -> None:
90+
idx = pd.period_range("2020-01-28 09:00", periods=4, freq="D")
91+
df = pd.DataFrame(data=4 * [range(2)], index=idx, columns=["a", "b"])
92+
93+
def f1(gb: ResamplerGroupBy):
94+
check(gb, PeriodIndexResamplerGroupby)
95+
96+
f1(df.groupby("a").resample("3min"))
97+
98+
99+
def test_natype() -> None:
100+
i64dt = pd.Int64Dtype()
101+
check(assert_type(i64dt.na_value, NAType), NAType)
102+
103+
104+
def test_nattype() -> None:
105+
td = pd.Timedelta("1 day")
106+
as_nat = pd.NaT
107+
108+
check(assert_type(td + as_nat, NaTType), NaTType)
109+
110+
111+
def test_expanding() -> None:
112+
df = pd.DataFrame({"B": [0, 1, 2, np.nan, 4]})
113+
114+
def f1(gb: Expanding):
115+
check(gb, Expanding)
116+
117+
f1(df.expanding())
118+
119+
120+
def test_expanding_groubpy() -> None:
121+
df = pd.DataFrame({"B": [0, 1, 2, np.nan, 4]})
122+
123+
def f1(gb: ExpandingGroupby):
124+
check(gb, ExpandingGroupby)
125+
126+
f1(df.groupby("B").expanding())
127+
128+
129+
def test_ewm() -> None:
130+
df = pd.DataFrame({"B": [0, 1, 2, np.nan, 4]})
131+
132+
def f1(gb: ExponentialMovingWindow):
133+
check(gb, ExponentialMovingWindow)
134+
135+
f1(df.ewm(2))
136+
137+
138+
def test_ewm_groubpy() -> None:
139+
df = pd.DataFrame({"B": [0, 1, 2, np.nan, 4]})
140+
141+
def f1(gb: ExponentialMovingWindowGroupby):
142+
check(gb, ExponentialMovingWindowGroupby)
143+
144+
f1(df.groupby("B").ewm(2))
145+
146+
147+
def test_json_reader() -> None:
148+
df = pd.DataFrame({"B": [0, 1, 2, np.nan, 4]})
149+
150+
def f1(gb: JsonReader):
151+
check(gb, JsonReader)
152+
153+
with ensure_clean() as path:
154+
check(assert_type(df.to_json(path), None), type(None))
155+
json_reader = read_json(path, chunksize=1, lines=True)
156+
f1(json_reader)
157+
json_reader.close()
158+
159+
160+
def test_resampler() -> None:
161+
s = pd.Series([1, 2, 3, 4, 5], index=pd.date_range("20130101", periods=5, freq="s"))
162+
163+
def f1(gb: Resampler):
164+
check(gb, Resampler)
165+
166+
f1(s.resample("3min"))
167+
168+
169+
def test_rolling() -> None:
170+
df = pd.DataFrame({"B": [0, 1, 2, np.nan, 4]})
171+
172+
def f1(gb: Rolling):
173+
check(gb, Rolling)
174+
175+
f1(df.rolling(2))
176+
177+
178+
def test_rolling_groupby() -> None:
179+
df = pd.DataFrame({"B": [0, 1, 2, np.nan, 4]})
180+
181+
def f1(gb: RollingGroupby):
182+
check(gb, RollingGroupby)
183+
184+
f1(df.groupby("B").rolling(2))
185+
186+
187+
def test_timegrouper() -> None:
188+
grouper = pd.Grouper(key="Publish date", freq="1W")
189+
190+
def f1(gb: TimeGrouper):
191+
check(gb, TimeGrouper)
192+
193+
f1(grouper)
194+
195+
196+
def test_window() -> None:
197+
ser = pd.Series([0, 1, 5, 2, 8])
198+
199+
def f1(gb: Window):
200+
check(gb, Window)
201+
202+
f1(ser.rolling(2, win_type="gaussian"))
203+
204+
205+
def test_statereader() -> None:
206+
df = pd.DataFrame([[1, 2], [3, 4]], columns=["col_1", "col_2"])
207+
time_stamp = pd.Timestamp(2000, 2, 29, 14, 21)
208+
variable_labels = {"col_1": "This is an example"}
209+
with ensure_clean() as path:
210+
df.to_stata(
211+
path, time_stamp=time_stamp, variable_labels=variable_labels, version=None
212+
)
213+
214+
def f1(gb: StataReader):
215+
check(gb, StataReader)
216+
217+
with StataReader(path) as reader:
218+
f1(reader)

0 commit comments

Comments
 (0)