Skip to content

assert types at runtime #114

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
merged 9 commits into from
Jul 10, 2022
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
9 changes: 9 additions & 0 deletions pandas-stubs/__init__.pyi
Original file line number Diff line number Diff line change
@@ -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,
Expand Down
21 changes: 21 additions & 0 deletions tests/__init__.py
Original file line number Diff line number Diff line change
@@ -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
52 changes: 31 additions & 21 deletions tests/test_frame.py
Original file line number Diff line number Diff line change
Expand Up @@ -20,6 +20,8 @@
import pytest
from typing_extensions import assert_type

from tests import check

from pandas.io.parsers import TextFileReader


Expand Down Expand Up @@ -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:
Expand Down Expand Up @@ -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:
Expand Down Expand Up @@ -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]
Expand Down Expand Up @@ -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"]
Expand Down Expand Up @@ -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:
Expand Down Expand Up @@ -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:
Expand Down
22 changes: 12 additions & 10 deletions tests/test_indexes.py
Original file line number Diff line number Diff line change
Expand Up @@ -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:
Expand All @@ -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)
47 changes: 25 additions & 22 deletions tests/test_interval.py
Original file line number Diff line number Diff line change
@@ -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")
Expand Down Expand Up @@ -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]
Loading