|
19 | 19 | import pytest |
20 | 20 | from typing_extensions import ( |
21 | 21 | Never, |
| 22 | + TypeAlias, |
22 | 23 | assert_type, |
23 | 24 | ) |
24 | 25 |
|
|
33 | 34 | pytest_warns_bounded, |
34 | 35 | ) |
35 | 36 |
|
| 37 | +if TYPE_CHECKING: |
| 38 | + from pandas.core.series import TimestampSeries |
| 39 | +else: |
| 40 | + TimestampSeries: TypeAlias = pd.Series |
| 41 | + |
36 | 42 |
|
37 | 43 | def test_types_to_datetime() -> None: |
38 | 44 | df = pd.DataFrame({"year": [2015, 2016], "month": [2, 3], "day": [4, 5]}) |
39 | | - r1: pd.Series = pd.to_datetime(df) |
| 45 | + check(assert_type(pd.to_datetime(df), TimestampSeries), pd.Series, pd.Timestamp) |
40 | 46 |
|
41 | | - r2: pd.Series = pd.to_datetime(df, unit="s", origin="unix") |
42 | | - r3: pd.Series = pd.to_datetime( |
43 | | - df, unit="ns", dayfirst=True, utc=False, format="%M:%D", exact=False |
| 47 | + check( |
| 48 | + assert_type(pd.to_datetime(df, unit="s", origin="unix"), TimestampSeries), |
| 49 | + pd.Series, |
| 50 | + pd.Timestamp, |
| 51 | + ) |
| 52 | + check( |
| 53 | + assert_type( |
| 54 | + pd.to_datetime( |
| 55 | + df, unit="ns", dayfirst=True, utc=False, format="%M:%D", exact=False |
| 56 | + ), |
| 57 | + TimestampSeries, |
| 58 | + ), |
| 59 | + pd.Series, |
| 60 | + pd.Timestamp, |
| 61 | + ) |
| 62 | + check( |
| 63 | + assert_type( |
| 64 | + pd.to_datetime([1, 2], unit="D", origin=pd.Timestamp("01/01/2000")), |
| 65 | + pd.DatetimeIndex, |
| 66 | + ), |
| 67 | + pd.DatetimeIndex, |
| 68 | + ) |
| 69 | + check( |
| 70 | + assert_type( |
| 71 | + pd.to_datetime([1, 2], unit="D", origin=3), |
| 72 | + pd.DatetimeIndex, |
| 73 | + ), |
| 74 | + pd.DatetimeIndex, |
| 75 | + ) |
| 76 | + check( |
| 77 | + assert_type( |
| 78 | + pd.to_datetime(["2022-01-03", "2022-02-22"]), |
| 79 | + pd.DatetimeIndex, |
| 80 | + ), |
| 81 | + pd.DatetimeIndex, |
44 | 82 | ) |
45 | | - r4: pd.DatetimeIndex = pd.to_datetime( |
46 | | - [1, 2], unit="D", origin=pd.Timestamp("01/01/2000") |
| 83 | + check( |
| 84 | + assert_type( |
| 85 | + pd.to_datetime(pd.Index(["2022-01-03", "2022-02-22"])), |
| 86 | + pd.DatetimeIndex, |
| 87 | + ), |
| 88 | + pd.DatetimeIndex, |
47 | 89 | ) |
48 | | - r5: pd.DatetimeIndex = pd.to_datetime([1, 2], unit="D", origin=3) |
49 | | - r6: pd.DatetimeIndex = pd.to_datetime(["2022-01-03", "2022-02-22"]) |
50 | | - r7: pd.DatetimeIndex = pd.to_datetime(pd.Index(["2022-01-03", "2022-02-22"])) |
51 | | - r8: pd.Series = pd.to_datetime( |
52 | | - {"year": [2015, 2016], "month": [2, 3], "day": [4, 5]} |
| 90 | + check( |
| 91 | + assert_type( |
| 92 | + pd.to_datetime({"year": [2015, 2016], "month": [2, 3], "day": [4, 5]}), |
| 93 | + TimestampSeries, |
| 94 | + ), |
| 95 | + pd.Series, |
| 96 | + pd.Timestamp, |
53 | 97 | ) |
54 | 98 |
|
55 | 99 |
|
@@ -179,7 +223,7 @@ def test_types_concat() -> None: |
179 | 223 |
|
180 | 224 | check( |
181 | 225 | assert_type( |
182 | | - pd.concat(map(lambda x: s2, ["some_value", 3]), axis=1), pd.DataFrame |
| 226 | + pd.concat(map(lambda _: s2, ["some_value", 3]), axis=1), pd.DataFrame |
183 | 227 | ), |
184 | 228 | pd.DataFrame, |
185 | 229 | ) |
@@ -306,14 +350,32 @@ def test_types_json_normalize() -> None: |
306 | 350 | {"name": {"given": "Mose", "family": "Regner"}}, |
307 | 351 | {"id": 2, "name": "Faye Raker"}, |
308 | 352 | ] |
309 | | - df1: pd.DataFrame = pd.json_normalize(data=data1) |
310 | | - df2: pd.DataFrame = pd.json_normalize(data=data1, max_level=0, sep=";") |
311 | | - df3: pd.DataFrame = pd.json_normalize( |
312 | | - data=data1, meta_prefix="id", record_prefix="name", errors="raise" |
| 353 | + check(assert_type(pd.json_normalize(data=data1), pd.DataFrame), pd.DataFrame) |
| 354 | + check( |
| 355 | + assert_type(pd.json_normalize(data=data1, max_level=0, sep=";"), pd.DataFrame), |
| 356 | + pd.DataFrame, |
| 357 | + ) |
| 358 | + check( |
| 359 | + assert_type( |
| 360 | + pd.json_normalize( |
| 361 | + data=data1, meta_prefix="id", record_prefix="name", errors="raise" |
| 362 | + ), |
| 363 | + pd.DataFrame, |
| 364 | + ), |
| 365 | + pd.DataFrame, |
| 366 | + ) |
| 367 | + check( |
| 368 | + assert_type( |
| 369 | + pd.json_normalize(data=data1, record_path=None, meta="id"), pd.DataFrame |
| 370 | + ), |
| 371 | + pd.DataFrame, |
313 | 372 | ) |
314 | | - df4: pd.DataFrame = pd.json_normalize(data=data1, record_path=None, meta="id") |
315 | 373 | data2: dict[str, Any] = {"name": {"given": "Mose", "family": "Regner"}} |
316 | | - df5: pd.DataFrame = pd.json_normalize(data=data2) |
| 374 | + check( |
| 375 | + assert_type(data2, dict[str, Any]), |
| 376 | + dict, |
| 377 | + ) |
| 378 | + check(assert_type(pd.json_normalize(data=data2), pd.DataFrame), pd.DataFrame) |
317 | 379 |
|
318 | 380 |
|
319 | 381 | def test_isna() -> None: |
|
0 commit comments