diff --git a/pandas-stubs/__init__.pyi b/pandas-stubs/__init__.pyi index 5b411f825..0810d8ac5 100644 --- a/pandas-stubs/__init__.pyi +++ b/pandas-stubs/__init__.pyi @@ -1,5 +1,14 @@ import pandas.testing as testing +from . import ( + api as api, + arrays as arrays, + errors as errors, + io as io, + plotting as plotting, + testing as testing, + tseries as tseries, +) from ._config import ( describe_option as describe_option, get_option as get_option, diff --git a/tests/__init__.py b/tests/__init__.py index e69de29bb..b4108beaf 100644 --- a/tests/__init__.py +++ b/tests/__init__.py @@ -0,0 +1,21 @@ +from __future__ import annotations + + +def check( + actual: object, klass: type, dtype: type | None = None, attr: str = "left" +) -> None: + + if not isinstance(actual, klass): + raise RuntimeError(f"Expected type '{klass}' but got '{type(actual)}'") + if dtype is None: + return None + + if hasattr(actual, "__iter__"): + value = next(iter(actual)) # type: ignore[call-overload] + else: + assert hasattr(actual, attr) + value = getattr(actual, attr) # type: ignore[attr-defined] + + if not isinstance(value, dtype): + raise RuntimeError(f"Expected type '{dtype}' but got '{type(value)}'") + return None diff --git a/tests/test_frame.py b/tests/test_frame.py index 16a0a95c5..a50950778 100644 --- a/tests/test_frame.py +++ b/tests/test_frame.py @@ -20,6 +20,8 @@ import pytest from typing_extensions import assert_type +from tests import check + from pandas.io.parsers import TextFileReader @@ -188,8 +190,8 @@ def test_types_assign() -> None: def test_types_sample() -> None: df = pd.DataFrame(data={"col1": [1, 2], "col2": [3, 4]}) # GH 67 - assert_type(df.sample(frac=0.5), pd.DataFrame) - assert_type(df.sample(n=1), pd.DataFrame) + check(assert_type(df.sample(frac=0.5), pd.DataFrame), pd.DataFrame) + check(assert_type(df.sample(n=1), pd.DataFrame), pd.DataFrame) def test_types_nlargest_nsmallest() -> None: @@ -576,10 +578,18 @@ def test_types_groupby_any() -> None: "col3": [False, False, False], } ) - assert_type(df.groupby("col1").any(), "pd.DataFrame") - assert_type(df.groupby("col1").all(), "pd.DataFrame") - assert_type(df.groupby("col1")["col2"].any(), "pd.Series[bool]") - assert_type(df.groupby("col1")["col2"].any(), "pd.Series[bool]") + check(assert_type(df.groupby("col1").any(), pd.DataFrame), pd.DataFrame) + check(assert_type(df.groupby("col1").all(), pd.DataFrame), pd.DataFrame) + check( + assert_type(df.groupby("col1")["col2"].any(), "pd.Series[bool]"), + pd.Series, + bool, + ) + check( + assert_type(df.groupby("col1")["col2"].any(), "pd.Series[bool]"), + pd.Series, + bool, + ) def test_types_merge() -> None: @@ -966,7 +976,7 @@ def test_read_csv() -> None: def test_groupby_series_methods() -> None: df = pd.DataFrame({"x": [1, 2, 2, 3, 3], "y": [10, 20, 30, 40, 50]}) gb = df.groupby("x")["y"] - assert_type(gb.describe(), "pd.DataFrame") + check(assert_type(gb.describe(), pd.DataFrame), pd.DataFrame) gb.count().loc[2] gb.pct_change().loc[2] gb.bfill().loc[2] @@ -1006,9 +1016,9 @@ def test_compute_values(): def test_sum_get_add() -> None: df = pd.DataFrame({"x": [1, 2, 3, 4, 5], "y": [10, 20, 30, 40, 50]}) s = df["x"] - assert_type(s, "pd.Series") + check(assert_type(s, pd.Series), pd.Series) summer = df.sum(axis=1) - assert_type(summer, "pd.Series") + check(assert_type(summer, pd.Series), pd.Series) s2: pd.Series = s + summer s3: pd.Series = s + df["y"] @@ -1036,7 +1046,7 @@ def test_getmultiindex_columns() -> None: def test_frame_getitem_isin() -> None: df = pd.DataFrame({"x": [1, 2, 3, 4, 5]}, index=[1, 2, 3, 4, 5]) - assert_type(df[df.index.isin([1, 3, 5])], "pd.DataFrame") + check(assert_type(df[df.index.isin([1, 3, 5])], pd.DataFrame), pd.DataFrame) def test_read_excel() -> None: @@ -1072,32 +1082,32 @@ def test_join() -> None: def test_types_ffill() -> None: # GH 44 df = pd.DataFrame([[1, 2, 3]]) - assert_type(df.ffill(), pd.DataFrame) - assert_type(df.ffill(inplace=False), pd.DataFrame) - assert_type(df.ffill(inplace=True), None) + check(assert_type(df.ffill(), pd.DataFrame), pd.DataFrame) + check(assert_type(df.ffill(inplace=False), pd.DataFrame), pd.DataFrame) + assert assert_type(df.ffill(inplace=True), None) is None def test_types_bfill() -> None: # GH 44 df = pd.DataFrame([[1, 2, 3]]) - assert_type(df.bfill(), pd.DataFrame) - assert_type(df.bfill(inplace=False), pd.DataFrame) - assert_type(df.bfill(inplace=True), None) + check(assert_type(df.bfill(), pd.DataFrame), pd.DataFrame) + check(assert_type(df.bfill(inplace=False), pd.DataFrame), pd.DataFrame) + assert assert_type(df.bfill(inplace=True), None) is None def test_types_replace() -> None: # GH 44 df = pd.DataFrame([[1, 2, 3]]) - assert_type(df.replace(1, 2), pd.DataFrame) - assert_type(df.replace(1, 2, inplace=False), pd.DataFrame) - assert_type(df.replace(1, 2, inplace=True), None) + check(assert_type(df.replace(1, 2), pd.DataFrame), pd.DataFrame) + check(assert_type(df.replace(1, 2, inplace=False), pd.DataFrame), pd.DataFrame) + assert assert_type(df.replace(1, 2, inplace=True), None) is None def test_loop_dataframe() -> None: # GH 70 df = pd.DataFrame({"x": [1, 2, 3]}) for c in df: - assert_type(df[c], pd.Series) + check(assert_type(df[c], pd.Series), pd.Series) def test_groupby_index() -> None: @@ -1105,7 +1115,7 @@ def test_groupby_index() -> None: df = pd.DataFrame( data={"col1": [1, 1, 2], "col2": [3, 4, 5], "col3": [0, 1, 0]} ).set_index("col1") - assert_type(df.groupby(df.index).min(), pd.DataFrame) + check(assert_type(df.groupby(df.index).min(), pd.DataFrame), pd.DataFrame) def test_iloc_npint() -> None: diff --git a/tests/test_indexes.py b/tests/test_indexes.py index bba230f0f..cc15513e3 100644 --- a/tests/test_indexes.py +++ b/tests/test_indexes.py @@ -3,42 +3,44 @@ import pandas as pd from typing_extensions import assert_type +from tests import check + def test_index_unique() -> None: df = pd.DataFrame({"x": [1, 2, 3, 4]}, index=pd.Index([1, 2, 3, 2])) ind = df.index - assert_type(ind, "pd.Index") + check(assert_type(ind, pd.Index), pd.Index) i2 = ind.unique() - assert_type(i2, "pd.Index") + check(assert_type(i2, pd.Index), pd.Index) def test_index_isin() -> None: ind = pd.Index([1, 2, 3, 4, 5]) isin = ind.isin([2, 4]) - assert_type(isin, npt.NDArray[np.bool_]) + check(assert_type(isin, npt.NDArray[np.bool_]), np.ndarray, np.bool_) def test_index_astype() -> None: indi = pd.Index([1, 2, 3]) inds = pd.Index(["a", "b", "c"]) indc = indi.astype(inds.dtype) - assert_type(indc, "pd.Index") + check(assert_type(indc, pd.Index), pd.Index) mi = pd.MultiIndex.from_product([["a", "b"], ["c", "d"]], names=["ab", "cd"]) mia = mi.astype(object) # object is only valid parameter for MultiIndex.astype() - assert_type(mia, "pd.MultiIndex") + check(assert_type(mia, pd.MultiIndex), pd.MultiIndex) def test_multiindex_get_level_values() -> None: mi = pd.MultiIndex.from_product([["a", "b"], ["c", "d"]], names=["ab", "cd"]) i1 = mi.get_level_values("ab") - assert_type(i1, "pd.Index") + check(assert_type(i1, pd.Index), pd.Index) def test_index_tolist() -> None: i1 = pd.Index([1, 2, 3]) - l1 = i1.tolist() - i2 = i1.to_list() + check(assert_type(i1.tolist(), list), list, int) + check(assert_type(i1.to_list(), list), list, int) def test_column_getitem() -> None: @@ -46,7 +48,7 @@ def test_column_getitem() -> None: df = pd.DataFrame([[1, 2, 3]], columns=["a", "b", "c"]) column = df.columns[0] - a = df[column] + check(assert_type(df[column], pd.Series), pd.Series, int) def test_column_contains() -> None: @@ -63,4 +65,4 @@ def test_column_contains() -> None: def test_difference_none() -> None: # https://github.com/pandas-dev/pandas-stubs/issues/17 ind = pd.Index([1, 2, 3]) - id = ind.difference([1, None]) + check(assert_type(ind.difference([1, None]), "pd.Index"), pd.Index, int) diff --git a/tests/test_interval.py b/tests/test_interval.py index edf8baa92..53973dd16 100644 --- a/tests/test_interval.py +++ b/tests/test_interval.py @@ -1,9 +1,12 @@ # flake8: noqa: F841 from typing import TYPE_CHECKING +import numpy as np import pandas as pd from typing_extensions import assert_type +from tests import check + def test_interval_init() -> None: i1: pd.Interval = pd.Interval(1, 2, closed="both") @@ -36,16 +39,16 @@ def test_interval_length() -> None: i1 = pd.Interval( pd.Timestamp("2000-01-01"), pd.Timestamp("2000-01-03"), closed="both" ) - assert_type(i1.length, "pd.Timedelta") - assert_type(i1.left, "pd.Timestamp") - assert_type(i1.right, "pd.Timestamp") - assert_type(i1.mid, "pd.Timestamp") + check(assert_type(i1.length, pd.Timedelta), pd.Timedelta) + check(assert_type(i1.left, pd.Timestamp), pd.Timestamp) + check(assert_type(i1.right, pd.Timestamp), pd.Timestamp) + check(assert_type(i1.mid, pd.Timestamp), pd.Timestamp) i1.length.total_seconds() inres = pd.Timestamp("2001-01-02") in i1 - assert_type(inres, "bool") + check(assert_type(inres, bool), bool) idres = i1 + pd.Timedelta(seconds=20) - assert_type(idres, "pd.Interval[pd.Timestamp]") + check(assert_type(idres, "pd.Interval[pd.Timestamp]"), pd.Interval, pd.Timestamp) if TYPE_CHECKING: 20 in i1 # type: ignore[operator] i1 + pd.Timestamp("2000-03-03") # type: ignore[operator] @@ -53,32 +56,32 @@ def test_interval_length() -> None: i1 * pd.Timedelta(seconds=20) # type: ignore[operator] i2 = pd.Interval(10, 20) - assert_type(i2.length, "int") - assert_type(i2.left, "int") - assert_type(i2.right, "int") - assert_type(i2.mid, "float") + check(assert_type(i2.length, int), int) + check(assert_type(i2.left, int), int) + check(assert_type(i2.right, int), int) + check(assert_type(i2.mid, float), float) i2inres = 15 in i2 - assert_type(i2inres, "bool") - assert_type(i2 + 3, "pd.Interval[int]") - assert_type(i2 + 3.2, "pd.Interval[float]") - assert_type(i2 * 4, "pd.Interval[int]") - assert_type(i2 * 4.2, "pd.Interval[float]") + check(assert_type(i2inres, bool), bool) + check(assert_type(i2 + 3, "pd.Interval[int]"), pd.Interval, int) + check(assert_type(i2 + 3.2, "pd.Interval[float]"), pd.Interval, float) + check(assert_type(i2 * 4, "pd.Interval[int]"), pd.Interval, int) + check(assert_type(i2 * 4.2, "pd.Interval[float]"), pd.Interval, float) if TYPE_CHECKING: pd.Timestamp("2001-01-02") in i2 # type: ignore[operator] i2 + pd.Timedelta(seconds=20) # type: ignore[operator] i3 = pd.Interval(13.2, 19.5) - assert_type(i3.length, "float") - assert_type(i3.left, "float") - assert_type(i3.right, "float") - assert_type(i3.mid, "float") + check(assert_type(i3.length, float), float) + check(assert_type(i3.left, float), float) + check(assert_type(i3.right, float), float) + check(assert_type(i3.mid, float), float) i3inres = 15.4 in i3 - assert_type(i3inres, "bool") - assert_type(i3 + 3, "pd.Interval[float]") - assert_type(i3 * 3, "pd.Interval[float]") + check(assert_type(i3inres, bool), bool) + check(assert_type(i3 + 3, "pd.Interval[float]"), pd.Interval, float) + check(assert_type(i3 * 3, "pd.Interval[float]"), pd.Interval, float) if TYPE_CHECKING: pd.Timestamp("2001-01-02") in i3 # type: ignore[operator] i3 + pd.Timedelta(seconds=20) # type: ignore[operator] diff --git a/tests/test_pandas.py b/tests/test_pandas.py index 58341c451..ad9b4dc05 100644 --- a/tests/test_pandas.py +++ b/tests/test_pandas.py @@ -13,6 +13,8 @@ import pytest from typing_extensions import assert_type +from tests import check + from pandas.io.parsers import TextFileReader @@ -40,12 +42,18 @@ def test_types_concat() -> None: s: pd.Series = pd.Series([0, 1, -10]) s2: pd.Series = pd.Series([7, -5, 10]) - assert_type(pd.concat([s, s2]), "pd.Series") - assert_type(pd.concat([s, s2], axis=1), "pd.DataFrame") - assert_type(pd.concat([s, s2], keys=["first", "second"], sort=True), "pd.Series") - assert_type( - pd.concat([s, s2], keys=["first", "second"], names=["source", "row"]), - "pd.Series", + check(assert_type(pd.concat([s, s2]), pd.Series), pd.Series) + check(assert_type(pd.concat([s, s2], axis=1), pd.DataFrame), pd.DataFrame) + check( + assert_type(pd.concat([s, s2], keys=["first", "second"], sort=True), pd.Series), + pd.Series, + ) + check( + assert_type( + pd.concat([s, s2], keys=["first", "second"], names=["source", "row"]), + pd.Series, + ), + pd.Series, ) # Depends on the axis @@ -59,14 +67,20 @@ def test_types_concat() -> None: df = pd.DataFrame(data={"col1": [1, 2], "col2": [3, 4]}) df2 = pd.DataFrame(data={"col1": [10, 20], "col2": [30, 40]}) - assert_type(pd.concat([df, df2]), "pd.DataFrame") - assert_type(pd.concat([df, df2], axis=1), "pd.DataFrame") - assert_type( - pd.concat([df, df2], keys=["first", "second"], sort=True), "pd.DataFrame" + check(assert_type(pd.concat([df, df2]), pd.DataFrame), pd.DataFrame) + check(assert_type(pd.concat([df, df2], axis=1), pd.DataFrame), pd.DataFrame) + check( + assert_type( + pd.concat([df, df2], keys=["first", "second"], sort=True), pd.DataFrame + ), + pd.DataFrame, ) - assert_type( - pd.concat([df, df2], keys=["first", "second"], names=["source", "row"]), - "pd.DataFrame", + check( + assert_type( + pd.concat([df, df2], keys=["first", "second"], names=["source", "row"]), + pd.DataFrame, + ), + pd.DataFrame, ) result: pd.DataFrame = pd.concat( @@ -149,10 +163,10 @@ def test_types_read_csv() -> None: def test_isna() -> None: s = pd.Series([1, np.nan, 3.2]) - assert_type(pd.isna(s), "pd.Series[bool]") + check(assert_type(pd.isna(s), "pd.Series[bool]"), pd.Series, bool) b: bool = pd.isna(np.nan) ar: np.ndarray = pd.isna(s.to_list()) - assert_type(pd.notna(s), "pd.Series[bool]") + check(assert_type(pd.notna(s), "pd.Series[bool]"), pd.Series, bool) b2: bool = pd.notna(np.nan) ar2: np.ndarray = pd.notna(s.to_list()) @@ -160,9 +174,12 @@ def test_isna() -> None: # GH 55 def test_read_xml() -> None: if TYPE_CHECKING: # Skip running pytest - assert_type( - pd.read_xml( - "path/to/file", xpath=".//row", stylesheet="path/to/stylesheet" + check( + assert_type( + pd.read_xml( + "path/to/file", xpath=".//row", stylesheet="path/to/stylesheet" + ), + pd.DataFrame, ), pd.DataFrame, ) diff --git a/tests/test_series.py b/tests/test_series.py index a35c16875..372ccfc07 100644 --- a/tests/test_series.py +++ b/tests/test_series.py @@ -17,6 +17,8 @@ from pandas._typing import Scalar +from tests import check + def test_types_init() -> None: pd.Series(1) @@ -38,15 +40,15 @@ def test_types_init() -> None: def test_types_any() -> None: - res1: bool = pd.Series([False, False]).any() - res2: bool = pd.Series([False, False]).any(bool_only=False) - res3: bool = pd.Series([np.nan]).any(skipna=False) + check(assert_type(pd.Series([False, False]).any(), bool), np.bool_) + check(assert_type(pd.Series([False, False]).any(bool_only=False), bool), np.bool_) + check(assert_type(pd.Series([np.nan]).any(skipna=False), bool), np.bool_) def test_types_all() -> None: - res1: bool = pd.Series([False, False]).all() - res2: bool = pd.Series([False, False]).all(bool_only=False) - res3: bool = pd.Series([np.nan]).all(skipna=False) + check(assert_type(pd.Series([False, False]).all(), bool), np.bool_) + check(assert_type(pd.Series([False, False]).all(bool_only=False), bool), np.bool_) + check(assert_type(pd.Series([np.nan]).all(skipna=False), bool), np.bool_) def test_types_csv() -> None: @@ -72,7 +74,7 @@ def test_types_csv() -> None: def test_types_copy() -> None: s = pd.Series(data=[1, 2, 3, 4]) - s2: pd.Series = s.copy() + check(assert_type(s.copy(), pd.Series), pd.Series, int) def test_types_select() -> None: @@ -101,8 +103,8 @@ def test_types_loc_at() -> None: def test_multiindex_loc() -> None: s = pd.Series([1, 2, 3, 4], index=pd.MultiIndex.from_product([[1, 2], ["a", "b"]])) - assert_type(s.loc[1, :], "pd.Series") - assert_type(s.loc[pd.Index([1]), :], "pd.Series") + check(assert_type(s.loc[1, :], pd.Series), pd.Series) + check(assert_type(s.loc[pd.Index([1]), :], pd.Series), pd.Series) def test_types_boolean_indexing() -> None: @@ -157,11 +159,11 @@ def test_types_setting() -> None: def test_types_drop() -> None: s = pd.Series([0, 1, 2]) - assert_type(s.drop(0), pd.Series) - assert_type(s.drop([0, 1]), pd.Series) - assert_type(s.drop(0, axis=0), pd.Series) - assert_type(s.drop([0, 1], inplace=True, errors="raise"), None) - assert_type(s.drop([0, 1], inplace=True, errors="ignore"), None) + check(assert_type(s.drop(0), pd.Series), pd.Series) + check(assert_type(s.drop([0, 1]), pd.Series), pd.Series) + check(assert_type(s.drop(0, axis=0), pd.Series), pd.Series) + assert assert_type(s.drop([0, 1], inplace=True, errors="raise"), None) is None + assert assert_type(s.drop([0, 1], inplace=True, errors="ignore"), None) is None def test_types_drop_multilevel() -> None: @@ -175,26 +177,26 @@ def test_types_drop_multilevel() -> None: def test_types_dropna() -> None: s = pd.Series([1, np.nan, np.nan]) - assert_type(s.dropna(), pd.Series) - assert_type(s.dropna(axis=0, inplace=True), None) + check(assert_type(s.dropna(), pd.Series), pd.Series) + assert assert_type(s.dropna(axis=0, inplace=True), None) is None def test_types_fillna() -> None: s = pd.Series([1, np.nan, np.nan, 3]) - assert_type(s.fillna(0), pd.Series) - assert_type(s.fillna(0, axis="index"), pd.Series) - assert_type(s.fillna(method="backfill", axis=0), pd.Series) - assert_type(s.fillna(method="bfill", inplace=True), None) - assert_type(s.fillna(method="pad"), pd.Series) - assert_type(s.fillna(method="ffill", limit=1), pd.Series) + check(assert_type(s.fillna(0), pd.Series), pd.Series) + check(assert_type(s.fillna(0, axis="index"), pd.Series), pd.Series) + check(assert_type(s.fillna(method="backfill", axis=0), pd.Series), pd.Series) + assert assert_type(s.fillna(method="bfill", inplace=True), None) is None + check(assert_type(s.fillna(method="pad"), pd.Series), pd.Series) + check(assert_type(s.fillna(method="ffill", limit=1), pd.Series), pd.Series) def test_types_sort_index() -> None: s = pd.Series([1, 2, 3], index=[2, 3, 1]) - assert_type(s.sort_index(), pd.Series) - assert_type(s.sort_index(ascending=False), pd.Series) - assert_type(s.sort_index(ascending=False, inplace=True), None) - assert_type(s.sort_index(kind="mergesort"), pd.Series) + check(assert_type(s.sort_index(), pd.Series), pd.Series) + check(assert_type(s.sort_index(ascending=False), pd.Series), pd.Series) + assert assert_type(s.sort_index(ascending=False, inplace=True), None) is None + check(assert_type(s.sort_index(kind="mergesort"), pd.Series), pd.Series) # This was added in 1.1.0 https://pandas.pydata.org/docs/whatsnew/v1.1.0.html @@ -205,13 +207,13 @@ def test_types_sort_index_with_key() -> None: def test_types_sort_values() -> None: s = pd.Series([4, 2, 1, 3]) - assert_type(s.sort_values(), pd.Series) + check(assert_type(s.sort_values(), pd.Series), pd.Series) with pytest.warns(FutureWarning, match="In a future version of pandas"): - assert_type(s.sort_values(0), pd.Series) - assert_type(s.sort_values(ascending=False), pd.Series) - assert_type(s.sort_values(inplace=True, kind="quicksort"), None) - assert_type(s.sort_values(na_position="last"), pd.Series) - assert_type(s.sort_values(ignore_index=True), pd.Series) + check(assert_type(s.sort_values(0), pd.Series), pd.Series) + check(assert_type(s.sort_values(ascending=False), pd.Series), pd.Series) + assert assert_type(s.sort_values(inplace=True, kind="quicksort"), None) is None + check(assert_type(s.sort_values(na_position="last"), pd.Series), pd.Series) + check(assert_type(s.sort_values(ignore_index=True), pd.Series), pd.Series) # This was added in 1.1.0 https://pandas.pydata.org/docs/whatsnew/v1.1.0.html @@ -549,25 +551,25 @@ def test_types_values() -> None: def test_types_rename() -> None: # Scalar s1 = pd.Series([1, 2, 3]).rename("A") - assert_type(s1, "pd.Series") + check(assert_type(s1, pd.Series), pd.Series) # Hashable Sequence s2 = pd.Series([1, 2, 3]).rename(("A", "B")) - assert_type(s2, "pd.Series") + check(assert_type(s2, pd.Series), pd.Series) # Optional s3 = pd.Series([1, 2, 3]).rename(None) - assert_type(s3, "pd.Series") + check(assert_type(s3, pd.Series), pd.Series) # Functions def add1(x: int) -> int: return x + 1 s4 = pd.Series([1, 2, 3]).rename(add1) - assert_type(s4, "pd.Series") + check(assert_type(s4, pd.Series), pd.Series) # Dictionary s5 = pd.Series([1, 2, 3]).rename({1: 10}) - assert_type(s5, "pd.Series") + check(assert_type(s5, pd.Series), pd.Series) # inplace s6: None = pd.Series([1, 2, 3]).rename("A", inplace=True) @@ -583,9 +585,9 @@ def test_types_ne() -> None: def test_types_bfill() -> None: s1 = pd.Series([1, 2, 3]) - assert_type(s1.bfill(), pd.Series) - assert_type(s1.bfill(inplace=False), pd.Series) - assert_type(s1.bfill(inplace=True), None) + check(assert_type(s1.bfill(), pd.Series), pd.Series) + check(assert_type(s1.bfill(inplace=False), pd.Series), pd.Series) + assert assert_type(s1.bfill(inplace=True), None) is None def test_types_ewm() -> None: @@ -600,9 +602,9 @@ def test_types_ewm() -> None: def test_types_ffill() -> None: s1 = pd.Series([1, 2, 3]) - assert_type(s1.ffill(), pd.Series) - assert_type(s1.ffill(inplace=False), pd.Series) - assert_type(s1.ffill(inplace=True), None) + check(assert_type(s1.ffill(), pd.Series), pd.Series) + check(assert_type(s1.ffill(inplace=False), pd.Series), pd.Series) + assert assert_type(s1.ffill(inplace=True), None) is None def test_types_as_type() -> None: @@ -637,10 +639,10 @@ def test_series_min_max_sub_axis() -> None: ss = s1 - s2 sm = s1 * s2 sd = s1 / s2 - assert_type(sa, "pd.Series") - assert_type(ss, "pd.Series") - assert_type(sm, "pd.Series") - assert_type(sd, "pd.Series") + check(assert_type(sa, pd.Series), pd.Series) + check(assert_type(ss, pd.Series), pd.Series) + check(assert_type(sm, pd.Series), pd.Series) + check(assert_type(sd, pd.Series), pd.Series) def test_series_index_isin() -> None: @@ -649,19 +651,19 @@ def test_series_index_isin() -> None: t2 = s.loc[~s.index.isin([1, 3])] t3 = s[s.index.isin([1, 3])] t4 = s[~s.index.isin([1, 3])] - assert_type(t1, "pd.Series") - assert_type(t2, "pd.Series") - assert_type(t3, "pd.Series") - assert_type(t4, "pd.Series") + check(assert_type(t1, pd.Series), pd.Series) + check(assert_type(t2, pd.Series), pd.Series) + check(assert_type(t3, pd.Series), pd.Series) + check(assert_type(t4, pd.Series), pd.Series) def test_series_invert() -> None: s1 = pd.Series([True, False, True]) s2 = ~s1 - assert_type(s2, "pd.Series[bool]") + check(assert_type(s2, "pd.Series[bool]"), pd.Series, bool) s3 = pd.Series([1, 2, 3]) - assert_type(s3[s2], "pd.Series") - assert_type(s3.loc[s2], "pd.Series") + check(assert_type(s3[s2], pd.Series), pd.Series) + check(assert_type(s3.loc[s2], pd.Series), pd.Series) def test_series_multiindex_getitem() -> None: @@ -674,13 +676,13 @@ def test_series_multiindex_getitem() -> None: def test_series_mul() -> None: s = pd.Series([1, 2, 3]) sm = s * 4 - assert_type(sm, "pd.Series") + check(assert_type(sm, pd.Series), pd.Series) ss = s - 4 - assert_type(ss, "pd.Series") + check(assert_type(ss, pd.Series), pd.Series) sm2 = s * s - assert_type(sm2, "pd.Series") + check(assert_type(sm2, pd.Series), pd.Series) sp = s + 4 - assert_type(sp, "pd.Series") + check(assert_type(sp, pd.Series), pd.Series) def test_reset_index() -> None: @@ -689,40 +691,40 @@ def test_reset_index() -> None: index=pd.MultiIndex.from_product([["a", "b"], ["c", "d"]], names=["ab", "cd"]), ) r1 = s.reset_index() - assert_type(r1, "pd.DataFrame") + check(assert_type(r1, pd.DataFrame), pd.DataFrame) r2 = s.reset_index(["ab"]) - assert_type(r2, "pd.DataFrame") + check(assert_type(r2, pd.DataFrame), pd.DataFrame) r3 = s.reset_index("ab") - assert_type(r3, "pd.DataFrame") + check(assert_type(r3, pd.DataFrame), pd.DataFrame) r4 = s.reset_index(drop=True) - assert_type(r4, "pd.Series") + check(assert_type(r4, pd.Series), pd.Series) r5 = s.reset_index(["ab"], drop=True) - assert_type(r5, "pd.Series") + check(assert_type(r5, pd.Series), pd.Series) def test_series_add_str() -> None: s = pd.Series(["abc", "def"]) - assert_type(s + "x", "pd.Series") - assert_type("x" + s, "pd.Series") + check(assert_type(s + "x", pd.Series), pd.Series) + check(assert_type("x" + s, pd.Series), pd.Series) def test_series_dtype() -> None: s = pd.Series(["abc", "def"], dtype=str) - assert_type(s, "pd.Series[str]") + check(assert_type(s, "pd.Series[str]"), pd.Series, str) def test_types_replace() -> None: # GH 44 s = pd.Series([1, 2, 3]) - assert_type(s.replace(1, 2), pd.Series) - assert_type(s.replace(1, 2, inplace=False), pd.Series) - assert_type(s.replace(1, 2, inplace=True), None) + check(assert_type(s.replace(1, 2), pd.Series), pd.Series) + check(assert_type(s.replace(1, 2, inplace=False), pd.Series), pd.Series) + assert assert_type(s.replace(1, 2, inplace=True), None) is None def test_cat_accessor() -> None: # GH 43 s = pd.Series(pd.Categorical(["a", "b", "a"], categories=["a", "b"])) - assert_type(s.cat.codes, "pd.Series[int]") + check(assert_type(s.cat.codes, "pd.Series[int]"), pd.Series, int) def test_cat_ctor_values() -> None: @@ -750,15 +752,15 @@ def test_iloc_getitem_ndarray() -> None: values_s = pd.Series(np.arange(10), name="a") - assert_type(values_s.iloc[indices_i8], "pd.Series") - assert_type(values_s.iloc[indices_i16], "pd.Series") - assert_type(values_s.iloc[indices_i32], "pd.Series") - assert_type(values_s.iloc[indices_i64], "pd.Series") + check(assert_type(values_s.iloc[indices_i8], pd.Series), pd.Series) + check(assert_type(values_s.iloc[indices_i16], pd.Series), pd.Series) + check(assert_type(values_s.iloc[indices_i32], pd.Series), pd.Series) + check(assert_type(values_s.iloc[indices_i64], pd.Series), pd.Series) - assert_type(values_s.iloc[indices_u8], "pd.Series") - assert_type(values_s.iloc[indices_u16], "pd.Series") - assert_type(values_s.iloc[indices_u32], "pd.Series") - assert_type(values_s.iloc[indices_u64], "pd.Series") + check(assert_type(values_s.iloc[indices_u8], pd.Series), pd.Series) + check(assert_type(values_s.iloc[indices_u16], pd.Series), pd.Series) + check(assert_type(values_s.iloc[indices_u32], pd.Series), pd.Series) + check(assert_type(values_s.iloc[indices_u64], pd.Series), pd.Series) def test_iloc_setitem_ndarray() -> None: diff --git a/tests/test_timefuncs.py b/tests/test_timefuncs.py index 8fb296675..465d667a9 100644 --- a/tests/test_timefuncs.py +++ b/tests/test_timefuncs.py @@ -14,6 +14,8 @@ from pandas._libs import NaTType +from tests import check + if TYPE_CHECKING: from pandas.core.series import ( TimedeltaSeries, @@ -66,9 +68,9 @@ def test_types_timestamp_series_comparisons() -> None: tssr = tss <= ts tssr2 = tss >= ts tssr3 = tss == ts - assert_type(tssr, "pd.Series[bool]") - assert_type(tssr2, "pd.Series[bool]") - assert_type(tssr3, "pd.Series[bool]") + check(assert_type(tssr, "pd.Series[bool]"), pd.Series, bool) + check(assert_type(tssr2, "pd.Series[bool]"), pd.Series, bool) + check(assert_type(tssr3, "pd.Series[bool]"), pd.Series, bool) def test_types_pydatetime() -> None: @@ -110,19 +112,19 @@ def test_timestamp_timedelta_series_arithmetic() -> None: td1 = pd.to_timedelta([2, 3], "seconds") ts2 = pd.to_datetime(pd.Series(["2022-03-08", "2022-03-10"])) r1 = ts1 - ts2 - assert_type(r1, "TimedeltaSeries") + check(assert_type(r1, "TimedeltaSeries"), pd.Series, pd.Timedelta) r2 = r1 / td1 - assert_type(r2, "pd.Series[float]") + check(assert_type(r2, "pd.Series[float]"), pd.Series, float) r3 = r1 - td1 - assert_type(r3, "TimedeltaSeries") + check(assert_type(r3, "TimedeltaSeries"), pd.Series, pd.Timedelta) r4 = pd.Timedelta(5, "days") / r1 - assert_type(r4, "pd.Series[float]") + check(assert_type(r4, "pd.Series[float]"), pd.Series, float) sb = pd.Series([1, 2]) == pd.Series([1, 3]) - assert_type(sb, "pd.Series[bool]") + check(assert_type(sb, "pd.Series[bool]"), pd.Series, bool) r5 = sb * r1 - assert_type(r5, "TimedeltaSeries") + check(assert_type(r5, "TimedeltaSeries"), pd.Series, pd.Timedelta) r6 = r1 * 4 - assert_type(r6, "TimedeltaSeries") + check(assert_type(r6, "TimedeltaSeries"), pd.Series, pd.Timedelta) def test_timestamp_dateoffset_arithmetic() -> None: @@ -136,16 +138,24 @@ def test_datetimeindex_plus_timedelta() -> None: dti = pd.to_datetime(["2022-03-08", "2022-03-15"]) td_s = pd.to_timedelta(pd.Series([10, 20]), "minutes") dti_td_s = dti + td_s - assert_type(dti_td_s, "TimestampSeries") + check( + assert_type(dti_td_s, "TimestampSeries"), + pd.Series, + pd.Timestamp, + ) td_dti_s = td_s + dti - assert_type(td_dti_s, "TimestampSeries") + check( + assert_type(td_dti_s, "TimestampSeries"), + pd.Series, + pd.Timestamp, + ) tdi = pd.to_timedelta([10, 20], "minutes") dti_tdi_dti = dti + tdi - assert_type(dti_tdi_dti, "pd.DatetimeIndex") + check(assert_type(dti_tdi_dti, "pd.DatetimeIndex"), pd.DatetimeIndex) tdi_dti_dti = tdi + dti - assert_type(tdi_dti_dti, "pd.DatetimeIndex") + check(assert_type(tdi_dti_dti, "pd.DatetimeIndex"), pd.DatetimeIndex) dti_td_dti = dti + pd.Timedelta(10, "minutes") - assert_type(dti_td_dti, "pd.DatetimeIndex") + check(assert_type(dti_td_dti, "pd.DatetimeIndex"), pd.DatetimeIndex) def test_timestamp_plus_timedelta_series() -> None: @@ -153,15 +163,17 @@ def test_timestamp_plus_timedelta_series() -> None: ts = pd.Timestamp("2022-03-05") td = pd.to_timedelta(pd.Series([10, 20]), "minutes") r3 = td + ts - assert_type(r3, "TimestampSeries") - # ignore type on next, because `tscheck` has Unknown dtype - assert_type(r3, "TimestampSeries") + check(assert_type(r3, "TimestampSeries"), pd.Series, pd.Timestamp) def test_timedelta_series_mult() -> None: df = pd.DataFrame({"x": [1, 3, 5], "y": [2, 2, 6]}) std = (df["x"] < df["y"]) * pd.Timedelta(10, "minutes") - assert_type(std, "TimedeltaSeries") + check( + assert_type(std, "TimedeltaSeries"), + pd.Series, + pd.Timedelta, + ) def test_timedelta_series_sum() -> None: @@ -193,45 +205,45 @@ def fail_on_adding_two_timestamps() -> None: def test_dtindex_tzinfo() -> None: # GH 71 dti = pd.date_range("2000-1-1", periods=10) - assert_type(dti.tzinfo, Optional[dt.tzinfo]) + assert assert_type(dti.tzinfo, Optional[dt.tzinfo]) is None def test_todatetime_fromnumpy() -> None: # GH 72 t1 = np.datetime64("2022-07-04 02:30") - assert_type(pd.to_datetime(t1), pd.Timestamp) + check(assert_type(pd.to_datetime(t1), pd.Timestamp), pd.Timestamp) def test_comparisons_datetimeindex() -> None: # GH 74 dti = pd.date_range("2000-01-01", "2000-01-10") ts = pd.Timestamp("2000-01-05") - assert_type((dti < ts), np_ndarray_bool) - assert_type((dti > ts), np_ndarray_bool) - assert_type((dti >= ts), np_ndarray_bool) - assert_type((dti <= ts), np_ndarray_bool) - assert_type((dti == ts), np_ndarray_bool) - assert_type((dti != ts), np_ndarray_bool) + check(assert_type((dti < ts), np_ndarray_bool), np.ndarray) + check(assert_type((dti > ts), np_ndarray_bool), np.ndarray) + check(assert_type((dti >= ts), np_ndarray_bool), np.ndarray) + check(assert_type((dti <= ts), np_ndarray_bool), np.ndarray) + check(assert_type((dti == ts), np_ndarray_bool), np.ndarray) + check(assert_type((dti != ts), np_ndarray_bool), np.ndarray) def test_to_datetime_nat() -> None: # GH 88 - assert isinstance( - assert_type(pd.to_datetime("2021-03-01", errors="ignore"), "pd.Timestamp"), + check( + assert_type(pd.to_datetime("2021-03-01", errors="ignore"), pd.Timestamp), pd.Timestamp, ) - assert isinstance( - assert_type(pd.to_datetime("2021-03-01", errors="raise"), "pd.Timestamp"), + check( + assert_type(pd.to_datetime("2021-03-01", errors="raise"), pd.Timestamp), pd.Timestamp, ) - assert isinstance( + check( assert_type( pd.to_datetime("2021-03-01", errors="coerce"), "Union[pd.Timestamp, NaTType]", ), pd.Timestamp, ) - assert isinstance( + check( assert_type( pd.to_datetime("not a date", errors="coerce"), "Union[pd.Timestamp, NaTType]",