From c72dbc748b2dcb8e6fffb8d51f8889d17d22f7d7 Mon Sep 17 00:00:00 2001 From: Amanda Bizzinotto Date: Tue, 6 Jun 2023 19:06:27 -0300 Subject: [PATCH 1/4] Update series.pyi and add test case --- pandas-stubs/core/series.pyi | 2 +- tests/test_series.py | 3 +++ 2 files changed, 4 insertions(+), 1 deletion(-) diff --git a/pandas-stubs/core/series.pyi b/pandas-stubs/core/series.pyi index ae6daa5da..17f508bb3 100644 --- a/pandas-stubs/core/series.pyi +++ b/pandas-stubs/core/series.pyi @@ -851,7 +851,7 @@ class Series(IndexOpsMixin, NDFrame, Generic[S1]): @overload def apply( self, - func: Callable[..., Scalar | Sequence | set | Mapping | None], + func: Callable[..., Scalar | Sequence | set | Mapping | NAType | None], convertDType: _bool = ..., args: tuple = ..., **kwds, diff --git a/tests/test_series.py b/tests/test_series.py index b34fd039f..7b2a1f817 100644 --- a/tests/test_series.py +++ b/tests/test_series.py @@ -35,6 +35,7 @@ ) import xarray as xr +from pandas._libs.missing import NAType from pandas._libs.tslibs.timedeltas import Timedelta from pandas._libs.tslibs.timestamps import Timestamp from pandas._typing import ( @@ -470,6 +471,8 @@ def get_depth(url: str) -> int: ss = s.astype(str) check(assert_type(ss.apply(get_depth), pd.Series), pd.Series, np.int64) + check(assert_type(s.apply(lambda x: pd.NA), pd.Series), pd.Series, NAType) + def test_types_element_wise_arithmetic() -> None: s = pd.Series([0, 1, -10]) From 99970d4f19f7cec37f21a23b87dbeaded5ab61a4 Mon Sep 17 00:00:00 2001 From: Amanda Bizzinotto Date: Mon, 12 Jun 2023 20:31:13 -0300 Subject: [PATCH 2/4] Add additional accepted types --- pandas-stubs/_libs/tslibs/period.pyi | 8 +++++++- pandas-stubs/core/indexes/period.pyi | 15 +++++++++++++-- 2 files changed, 20 insertions(+), 3 deletions(-) diff --git a/pandas-stubs/_libs/tslibs/period.pyi b/pandas-stubs/_libs/tslibs/period.pyi index d38dbcfc2..c41733ab0 100644 --- a/pandas-stubs/_libs/tslibs/period.pyi +++ b/pandas-stubs/_libs/tslibs/period.pyi @@ -5,6 +5,7 @@ from typing import ( ) import numpy as np +import pandas from pandas import ( Index, PeriodIndex, @@ -61,7 +62,12 @@ class PeriodMixin: class Period(PeriodMixin): def __init__( self, - value: Period | str | None = ..., + value: Period + | str + | datetime.datetime + | datetime.date + | Timestamp + | None = ..., freq: str | BaseOffset | None = ..., ordinal: int | None = ..., year: int | None = ..., diff --git a/pandas-stubs/core/indexes/period.pyi b/pandas-stubs/core/indexes/period.pyi index 8edd2cae7..49b23b102 100644 --- a/pandas-stubs/core/indexes/period.pyi +++ b/pandas-stubs/core/indexes/period.pyi @@ -1,4 +1,5 @@ from collections.abc import Hashable +import datetime from typing import overload import numpy as np @@ -84,8 +85,18 @@ class PeriodIndex( # type: ignore[misc] def freqstr(self) -> str: ... def period_range( - start: str | pd.Period | None = ..., - end: str | pd.Period | None = ..., + start: str + | datetime.datetime + | datetime.date + | pd.Timestamp + | pd.Period + | None = ..., + end: str + | datetime.datetime + | datetime.date + | pd.Timestamp + | pd.Period + | None = ..., periods: int | None = ..., freq: str | BaseOffset | None = ..., name: Hashable | None = ..., From ee46a6badd63c8f064a9aa57532c98ad7434282d Mon Sep 17 00:00:00 2001 From: Amanda Bizzinotto Date: Mon, 12 Jun 2023 20:47:12 -0300 Subject: [PATCH 3/4] Add test cases for additional types --- tests/test_scalars.py | 15 +++++++++++++++ tests/test_timefuncs.py | 21 +++++++++++++++++++++ 2 files changed, 36 insertions(+) diff --git a/tests/test_scalars.py b/tests/test_scalars.py index 510536502..958a03b61 100644 --- a/tests/test_scalars.py +++ b/tests/test_scalars.py @@ -1,5 +1,6 @@ from __future__ import annotations +import datetime import datetime as dt from typing import ( TYPE_CHECKING, @@ -1690,6 +1691,20 @@ def test_period_construction() -> None: pd.Period, ) check(assert_type(pd.Period(freq="Q", year=2012, quarter=2), pd.Period), pd.Period) + check( + assert_type( + pd.Period(value=datetime.datetime(2012, 1, 1), freq="D"), pd.Period + ), + pd.Period, + ) + check( + assert_type(pd.Period(value=datetime.date(2012, 1, 1), freq="D"), pd.Period), + pd.Period, + ) + check( + assert_type(pd.Period(value=pd.Timestamp(2012, 1, 1), freq="D"), pd.Period), + pd.Period, + ) def test_period_properties() -> None: diff --git a/tests/test_timefuncs.py b/tests/test_timefuncs.py index 01d6a29b4..8d1ef3309 100644 --- a/tests/test_timefuncs.py +++ b/tests/test_timefuncs.py @@ -838,6 +838,27 @@ def test_period_range() -> None: ), pd.PeriodIndex, ) + check( + assert_type( + pd.period_range(pd.Timestamp("2001-01-01"), end="2002-01-01", freq="Q"), + pd.PeriodIndex, + ), + pd.PeriodIndex, + ) + check( + assert_type( + pd.period_range(dt.datetime(2001, 1, 1), end="2002-01-01", freq="Q"), + pd.PeriodIndex, + ), + pd.PeriodIndex, + ) + check( + assert_type( + pd.period_range(dt.date(2001, 1, 1), end="2002-01-01", freq="Q"), + pd.PeriodIndex, + ), + pd.PeriodIndex, + ) def test_to_datetime_scalar() -> None: From 52719fbf079198862686f6566f3459ef0c7c8fb8 Mon Sep 17 00:00:00 2001 From: Amanda Bizzinotto Date: Tue, 13 Jun 2023 14:13:12 -0300 Subject: [PATCH 4/4] Fix tests and remove unnecessary import --- pandas-stubs/_libs/tslibs/period.pyi | 1 - tests/test_timefuncs.py | 10 +++++++--- 2 files changed, 7 insertions(+), 4 deletions(-) diff --git a/pandas-stubs/_libs/tslibs/period.pyi b/pandas-stubs/_libs/tslibs/period.pyi index c41733ab0..315030954 100644 --- a/pandas-stubs/_libs/tslibs/period.pyi +++ b/pandas-stubs/_libs/tslibs/period.pyi @@ -5,7 +5,6 @@ from typing import ( ) import numpy as np -import pandas from pandas import ( Index, PeriodIndex, diff --git a/tests/test_timefuncs.py b/tests/test_timefuncs.py index 8d1ef3309..20a965545 100644 --- a/tests/test_timefuncs.py +++ b/tests/test_timefuncs.py @@ -840,21 +840,25 @@ def test_period_range() -> None: ) check( assert_type( - pd.period_range(pd.Timestamp("2001-01-01"), end="2002-01-01", freq="Q"), + pd.period_range( + pd.Timestamp("2001-01-01"), end=pd.Timestamp("2002-01-01"), freq="Q" + ), pd.PeriodIndex, ), pd.PeriodIndex, ) check( assert_type( - pd.period_range(dt.datetime(2001, 1, 1), end="2002-01-01", freq="Q"), + pd.period_range( + dt.datetime(2001, 1, 1), end=dt.datetime(2002, 1, 1), freq="Q" + ), pd.PeriodIndex, ), pd.PeriodIndex, ) check( assert_type( - pd.period_range(dt.date(2001, 1, 1), end="2002-01-01", freq="Q"), + pd.period_range(dt.date(2001, 1, 1), end=dt.date(2002, 1, 1), freq="Q"), pd.PeriodIndex, ), pd.PeriodIndex,