Skip to content

Commit 9fcaba3

Browse files
MarcoGorelliMarcoGorelli
committed
api test, typing
1 parent 3d17e56 commit 9fcaba3

File tree

2 files changed

+80
-72
lines changed

2 files changed

+80
-72
lines changed

‎pandas/core/col.py

Lines changed: 79 additions & 72 deletions
Original file line numberDiff line numberDiff line change
@@ -15,22 +15,30 @@
1515
from pandas import DataFrame
1616

1717

18-
def parse_args(df: DataFrame, *args) -> tuple[Series]:
19-
return tuple([x(df) if isinstance(x, Expr) else x for x in args])
18+
def parse_args(df: DataFrame, *args: Any) -> tuple[Series]:
19+
return tuple([x._func(df) if isinstance(x, Expr) else x for x in args])
2020

2121

22-
def parse_kwargs(df: DataFrame, **kwargs) -> dict[Hashable, Series]:
22+
def parse_kwargs(df: DataFrame, **kwargs: Any) -> dict[Hashable, Series]:
2323
return {
24-
key: val(df) if isinstance(val, Expr) else val for key, val in kwargs.items()
24+
key: val._func(df) if isinstance(val, Expr) else val
25+
for key, val in kwargs.items()
2526
}
2627

2728

2829
class Expr:
29-
def __init__(self, func: Callable[[DataFrame], Series]) -> None:
30+
def __init__(self, func: Callable[[DataFrame], Any]) -> None:
3031
self._func = func
3132

3233
def __call__(self, df: DataFrame) -> Series:
33-
return self._func(df)
34+
result = self._func(df)
35+
if not isinstance(result, Series):
36+
msg = (
37+
"Expected function which returns Series, "
38+
f"got function which returns: {type(result)}"
39+
)
40+
raise TypeError(msg)
41+
return result
3442

3543
# namespaces
3644
@property
@@ -59,120 +67,119 @@ def struct(self) -> NamespaceExpr:
5967

6068
# Binary ops
6169

62-
def __add__(self, other) -> Expr:
70+
def __add__(self, other: Any) -> Expr:
6371
if isinstance(other, Expr):
64-
return Expr(lambda df: self(df).__add__(other(df)))
65-
return Expr(lambda df: self(df).__add__(other))
72+
return Expr(lambda df: self._func(df).__add__(other._func(df)))
73+
return Expr(lambda df: self._func(df).__add__(other))
6674

67-
def __radd__(self, other) -> Expr:
75+
def __radd__(self, other: Any) -> Expr:
6876
if isinstance(other, Expr):
69-
return Expr(lambda df: self(df).__radd__(other(df)))
70-
return Expr(lambda df: self(df).__radd__(other))
77+
return Expr(lambda df: self._func(df).__radd__(other._func(df)))
78+
return Expr(lambda df: self._func(df).__radd__(other))
7179

72-
def __sub__(self, other) -> Expr:
80+
def __sub__(self, other: Any) -> Expr:
7381
if isinstance(other, Expr):
74-
return Expr(lambda df: self(df).__sub__(other(df)))
75-
return Expr(lambda df: self(df).__sub__(other))
82+
return Expr(lambda df: self._func(df).__sub__(other._func(df)))
83+
return Expr(lambda df: self._func(df).__sub__(other))
7684

77-
def __rsub__(self, other) -> Expr:
85+
def __rsub__(self, other: Any) -> Expr:
7886
if isinstance(other, Expr):
79-
return Expr(lambda df: self(df).__rsub__(other(df)))
80-
return Expr(lambda df: self(df).__rsub__(other))
87+
return Expr(lambda df: self._func(df).__rsub__(other._func(df)))
88+
return Expr(lambda df: self._func(df).__rsub__(other))
8189

82-
def __mul__(self, other) -> Expr:
90+
def __mul__(self, other: Any) -> Expr:
8391
if isinstance(other, Expr):
84-
return Expr(lambda df: self(df).__mul__(other(df)))
85-
return Expr(lambda df: self(df).__mul__(other))
92+
return Expr(lambda df: self._func(df).__mul__(other._func(df)))
93+
return Expr(lambda df: self._func(df).__mul__(other))
8694

87-
def __rmul__(self, other) -> Expr:
95+
def __rmul__(self, other: Any) -> Expr:
8896
if isinstance(other, Expr):
89-
return Expr(lambda df: self(df).__rmul__(other(df)))
90-
return Expr(lambda df: self(df).__rmul__(other))
97+
return Expr(lambda df: self._func(df).__rmul__(other._func(df)))
98+
return Expr(lambda df: self._func(df).__rmul__(other))
9199

92-
def __truediv__(self, other) -> Expr:
100+
def __truediv__(self, other: Any) -> Expr:
93101
if isinstance(other, Expr):
94-
return Expr(lambda df: self(df).__truediv__(other(df)))
95-
return Expr(lambda df: self(df).__truediv__(other))
102+
return Expr(lambda df: self._func(df).__truediv__(other._func(df)))
103+
return Expr(lambda df: self._func(df).__truediv__(other))
96104

97-
def __rtruediv__(self, other) -> Expr:
105+
def __rtruediv__(self, other: Any) -> Expr:
98106
if isinstance(other, Expr):
99-
return Expr(lambda df: self(df).__rtruediv__(other(df)))
100-
return Expr(lambda df: self(df).__rtruediv__(other))
107+
return Expr(lambda df: self._func(df).__rtruediv__(other._func(df)))
108+
return Expr(lambda df: self._func(df).__rtruediv__(other))
101109

102-
def __floordiv__(self, other) -> Expr:
110+
def __floordiv__(self, other: Any) -> Expr:
103111
if isinstance(other, Expr):
104-
return Expr(lambda df: self(df).__floordiv__(other(df)))
105-
return Expr(lambda df: self(df).__floordiv__(other))
112+
return Expr(lambda df: self._func(df).__floordiv__(other._func(df)))
113+
return Expr(lambda df: self._func(df).__floordiv__(other))
106114

107-
def __rfloordiv__(self, other) -> Expr:
115+
def __rfloordiv__(self, other: Any) -> Expr:
108116
if isinstance(other, Expr):
109-
return Expr(lambda df: self(df).__rfloordiv__(other(df)))
110-
return Expr(lambda df: self(df).__rfloordiv__(other))
117+
return Expr(lambda df: self._func(df).__rfloordiv__(other._func(df)))
118+
return Expr(lambda df: self._func(df).__rfloordiv__(other))
111119

112-
def __ge__(self, other) -> Expr:
120+
def __ge__(self, other: Any) -> Expr:
113121
if isinstance(other, Expr):
114-
return Expr(lambda df: self(df).__ge__(other(df)))
115-
return Expr(lambda df: self(df).__ge__(other))
122+
return Expr(lambda df: self._func(df).__ge__(other._func(df)))
123+
return Expr(lambda df: self._func(df).__ge__(other))
116124

117-
def __gt__(self, other) -> Expr:
125+
def __gt__(self, other: Any) -> Expr:
118126
if isinstance(other, Expr):
119-
return Expr(lambda df: self(df).__gt__(other(df)))
120-
return Expr(lambda df: self(df).__gt__(other))
127+
return Expr(lambda df: self._func(df).__gt__(other._func(df)))
128+
return Expr(lambda df: self._func(df).__gt__(other))
121129

122-
def __le__(self, other) -> Expr:
130+
def __le__(self, other: Any) -> Expr:
123131
if isinstance(other, Expr):
124-
return Expr(lambda df: self(df).__le__(other(df)))
125-
return Expr(lambda df: self(df).__le__(other))
132+
return Expr(lambda df: self._func(df).__le__(other._func(df)))
133+
return Expr(lambda df: self._func(df).__le__(other))
126134

127-
def __lt__(self, other) -> Expr:
135+
def __lt__(self, other: Any) -> Expr:
128136
if isinstance(other, Expr):
129-
return Expr(lambda df: self(df).__lt__(other(df)))
130-
return Expr(lambda df: self(df).__lt__(other))
137+
return Expr(lambda df: self._func(df).__lt__(other._func(df)))
138+
return Expr(lambda df: self._func(df).__lt__(other))
131139

132-
def __eq__(self, other) -> Expr:
140+
def __eq__(self, other: object) -> Expr: # type: ignore[override]
133141
if isinstance(other, Expr):
134-
return Expr(lambda df: self(df).__eq__(other(df)))
135-
return Expr(lambda df: self(df).__eq__(other))
142+
return Expr(lambda df: self._func(df).__eq__(other._func(df)))
143+
return Expr(lambda df: self._func(df).__eq__(other))
136144

137-
def __neq__(self, other) -> Expr:
145+
def __ne__(self, other: object) -> Expr: # type: ignore[override]
138146
if isinstance(other, Expr):
139-
return Expr(lambda df: self(df).__neq__(other(df)))
140-
return Expr(lambda df: self(df).__neq__(other))
147+
return Expr(lambda df: self._func(df).__ne__(other._func(df)))
148+
return Expr(lambda df: self._func(df).__ne__(other))
141149

142-
def __mod__(self, other) -> Expr:
150+
def __mod__(self, other: Any) -> Expr:
143151
if isinstance(other, Expr):
144-
return Expr(lambda df: self(df).__mod__(other(df)))
145-
return Expr(lambda df: self(df).__mod__(other))
152+
return Expr(lambda df: self._func(df).__mod__(other._func(df)))
153+
return Expr(lambda df: self._func(df).__mod__(other))
146154

147155
# Everything else
148156

149-
def __getattr__(self, attr: str) -> Expr:
150-
def func(df: DataFrame, *args: Any, **kwargs: Any) -> Series:
151-
args = parse_args(df, *args)
152-
kwargs = parse_kwargs(df, **kwargs)
153-
return getattr(self(df), attr)(*args, **kwargs)
157+
# Function "pandas.core.col.Expr.str" is not valid as a type
158+
def __getattr__(self, attr: str, /) -> Any: # type: ignore[valid-type]
159+
def func(df: DataFrame, *args: Any, **kwargs: Any) -> Any:
160+
parsed_args = parse_args(df, *args)
161+
parsed_kwargs = parse_kwargs(df, **kwargs)
162+
return getattr(self(df), attr)(*parsed_args, **parsed_kwargs)
154163

155164
return lambda *args, **kwargs: Expr(lambda df: func(df, *args, **kwargs))
156165

157166

158167
class NamespaceExpr:
159-
def __init__(self, func: Callable[[DataFrame], Series], namespace: str) -> None:
168+
def __init__(self, func: Callable[[DataFrame], Any], namespace: str) -> None:
160169
self._func = func
161170
self._namespace = namespace
162171

163172
def __getattr__(self, attr: str) -> Any:
164173
if isinstance(getattr(getattr(Series, self._namespace), attr), property):
174+
return Expr(
175+
lambda df: getattr(getattr(self._func(df), self._namespace), attr)
176+
)
165177

166-
def func(df):
167-
return getattr(getattr(self._func(df), self._namespace), attr)
168-
169-
return Expr(func)
170-
171-
def func(df, *args, **kwargs):
172-
args = parse_args(df, *args)
173-
kwargs = parse_kwargs(df, **kwargs)
178+
def func(df: DataFrame, *args: Any, **kwargs: Any) -> Any:
179+
parsed_args = parse_args(df, *args)
180+
parsed_kwargs = parse_kwargs(df, **kwargs)
174181
return getattr(getattr(self._func(df), self._namespace), attr)(
175-
*args, **kwargs
182+
*parsed_args, **parsed_kwargs
176183
)
177184

178185
return lambda *args, **kwargs: Expr(lambda df: func(df, *args, **kwargs))

‎pandas/tests/api/test_api.py

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -107,6 +107,7 @@ class TestPDApi(Base):
107107
funcs = [
108108
"array",
109109
"bdate_range",
110+
"col",
110111
"concat",
111112
"crosstab",
112113
"cut",

0 commit comments

Comments
 (0)