assert types at runtime (#114)

* assert types at runtime

* use classes isntead of strings where possible

* check dtype

* check()

* unused imports

* check a few unused variables

* attr

Author: twoertwein

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,

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

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,40 +1082,40 @@ 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:
     # GH 42
     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:

tests/test_indexes.py

@@ -3,50 +3,52 @@
 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:
     # https://github.com/microsoft/python-type-stubs/issues/199#issuecomment-1132806594
     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)

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,49 +39,49 @@ 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]
         i1 * 3  # type: ignore[operator]
         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]