1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
|
"""Tests for pandas series and dataframe method registration."""
import pandas_flavor as pf
import pandas as pd
from pandas.core.groupby.generic import DataFrameGroupBy
def test_register_dataframe_method():
"""Test register_dataframe_method."""
@pf.register_dataframe_method
def dummy_func(df: pd.DataFrame) -> pd.DataFrame:
"""Dummy function.
Args:
df: a pandas DataFrame
Returns:
df: A pandas DataFrame.
"""
return df
df = pd.DataFrame()
df.dummy_func()
def test_register_series_method():
"""Test register_series_method."""
@pf.register_series_method
def dummy_func(s: pd.Series) -> pd.Series:
"""Dummy func.
Args:
s: A pandas Series.
Returns:
s: A pandas Series.
"""
return s
ser = pd.Series()
ser.dummy_func()
def test_register_groupby_method():
"""Test register_groupby_method."""
@pf.register_groupby_method
def dummy_func(by: DataFrameGroupBy) -> DataFrameGroupBy:
"""Dummy func.
Args:
by: A DataFrameGroupBy object.
Returns:
DataFrameGroupBy.
"""
return by
df = pd.DataFrame(
{
"Animal": ["Falcon"],
"Max Speed": [380.0],
}
)
by = df.groupby("Animal")
by.dummy_func()
|
