Skip to content

Commit 33665e2

Browse files
MarcoGorelliMarcoGorelli
committed
typing
1 parent 9fcaba3 commit 33665e2

File tree

1 file changed

+47
-79
lines changed

1 file changed

+47
-79
lines changed

‎pandas/core/col.py

Lines changed: 47 additions & 79 deletions
Original file line numberDiff line numberDiff line change
@@ -40,129 +40,97 @@ def __call__(self, df: DataFrame) -> Series:
4040
raise TypeError(msg)
4141
return result
4242

43-
# namespaces
44-
@property
45-
def dt(self) -> NamespaceExpr:
46-
return NamespaceExpr(self, "dt")
47-
48-
@property
49-
def str(self) -> NamespaceExpr:
50-
return NamespaceExpr(self, "str")
51-
52-
@property
53-
def cat(self) -> NamespaceExpr:
54-
return NamespaceExpr(self, "cat")
55-
56-
@property
57-
def list(self) -> NamespaceExpr:
58-
return NamespaceExpr(self, "list")
59-
60-
@property
61-
def sparse(self) -> NamespaceExpr:
62-
return NamespaceExpr(self, "sparse")
63-
64-
@property
65-
def struct(self) -> NamespaceExpr:
66-
return NamespaceExpr(self, "struct")
43+
def _with_binary_op(self, op: str, other: Any) -> Expr:
44+
if isinstance(other, Expr):
45+
return Expr(lambda df: getattr(self._func(df), op)(other._func(df)))
46+
return Expr(lambda df: getattr(self._func(df), op)(other))
6747

6848
# Binary ops
69-
7049
def __add__(self, other: Any) -> Expr:
71-
if isinstance(other, Expr):
72-
return Expr(lambda df: self._func(df).__add__(other._func(df)))
73-
return Expr(lambda df: self._func(df).__add__(other))
50+
return self._with_binary_op("__add__", other)
7451

7552
def __radd__(self, other: Any) -> Expr:
76-
if isinstance(other, Expr):
77-
return Expr(lambda df: self._func(df).__radd__(other._func(df)))
78-
return Expr(lambda df: self._func(df).__radd__(other))
53+
return self._with_binary_op("__radd__", other)
7954

8055
def __sub__(self, other: Any) -> Expr:
81-
if isinstance(other, Expr):
82-
return Expr(lambda df: self._func(df).__sub__(other._func(df)))
83-
return Expr(lambda df: self._func(df).__sub__(other))
56+
return self._with_binary_op("__sub__", other)
8457

8558
def __rsub__(self, other: Any) -> Expr:
86-
if isinstance(other, Expr):
87-
return Expr(lambda df: self._func(df).__rsub__(other._func(df)))
88-
return Expr(lambda df: self._func(df).__rsub__(other))
59+
return self._with_binary_op("__rsub__", other)
8960

9061
def __mul__(self, other: Any) -> Expr:
91-
if isinstance(other, Expr):
92-
return Expr(lambda df: self._func(df).__mul__(other._func(df)))
93-
return Expr(lambda df: self._func(df).__mul__(other))
62+
return self._with_binary_op("__mul__", other)
9463

9564
def __rmul__(self, other: Any) -> Expr:
96-
if isinstance(other, Expr):
97-
return Expr(lambda df: self._func(df).__rmul__(other._func(df)))
98-
return Expr(lambda df: self._func(df).__rmul__(other))
65+
return self._with_binary_op("__rmul__", other)
9966

10067
def __truediv__(self, other: Any) -> Expr:
101-
if isinstance(other, Expr):
102-
return Expr(lambda df: self._func(df).__truediv__(other._func(df)))
103-
return Expr(lambda df: self._func(df).__truediv__(other))
68+
return self._with_binary_op("__truediv__", other)
10469

10570
def __rtruediv__(self, other: Any) -> Expr:
106-
if isinstance(other, Expr):
107-
return Expr(lambda df: self._func(df).__rtruediv__(other._func(df)))
108-
return Expr(lambda df: self._func(df).__rtruediv__(other))
71+
return self._with_binary_op("__rtruediv__", other)
10972

11073
def __floordiv__(self, other: Any) -> Expr:
111-
if isinstance(other, Expr):
112-
return Expr(lambda df: self._func(df).__floordiv__(other._func(df)))
113-
return Expr(lambda df: self._func(df).__floordiv__(other))
74+
return self._with_binary_op("__floordiv__", other)
11475

11576
def __rfloordiv__(self, other: Any) -> Expr:
116-
if isinstance(other, Expr):
117-
return Expr(lambda df: self._func(df).__rfloordiv__(other._func(df)))
118-
return Expr(lambda df: self._func(df).__rfloordiv__(other))
77+
return self._with_binary_op("__rfloordiv__", other)
11978

12079
def __ge__(self, other: Any) -> Expr:
121-
if isinstance(other, Expr):
122-
return Expr(lambda df: self._func(df).__ge__(other._func(df)))
123-
return Expr(lambda df: self._func(df).__ge__(other))
80+
return self._with_binary_op("__ge__", other)
12481

12582
def __gt__(self, other: Any) -> Expr:
126-
if isinstance(other, Expr):
127-
return Expr(lambda df: self._func(df).__gt__(other._func(df)))
128-
return Expr(lambda df: self._func(df).__gt__(other))
83+
return self._with_binary_op("__gt__", other)
12984

13085
def __le__(self, other: Any) -> Expr:
131-
if isinstance(other, Expr):
132-
return Expr(lambda df: self._func(df).__le__(other._func(df)))
133-
return Expr(lambda df: self._func(df).__le__(other))
86+
return self._with_binary_op("__le__", other)
13487

13588
def __lt__(self, other: Any) -> Expr:
136-
if isinstance(other, Expr):
137-
return Expr(lambda df: self._func(df).__lt__(other._func(df)))
138-
return Expr(lambda df: self._func(df).__lt__(other))
89+
return self._with_binary_op("__lt__", other)
13990

14091
def __eq__(self, other: object) -> Expr: # type: ignore[override]
141-
if isinstance(other, Expr):
142-
return Expr(lambda df: self._func(df).__eq__(other._func(df)))
143-
return Expr(lambda df: self._func(df).__eq__(other))
92+
return self._with_binary_op("__eq__", other)
14493

14594
def __ne__(self, other: object) -> Expr: # type: ignore[override]
146-
if isinstance(other, Expr):
147-
return Expr(lambda df: self._func(df).__ne__(other._func(df)))
148-
return Expr(lambda df: self._func(df).__ne__(other))
95+
return self._with_binary_op("__ne__", other)
14996

15097
def __mod__(self, other: Any) -> Expr:
151-
if isinstance(other, Expr):
152-
return Expr(lambda df: self._func(df).__mod__(other._func(df)))
153-
return Expr(lambda df: self._func(df).__mod__(other))
98+
return self._with_binary_op("__mod__", other)
15499

155100
# Everything else
156-
157-
# Function "pandas.core.col.Expr.str" is not valid as a type
158-
def __getattr__(self, attr: str, /) -> Any: # type: ignore[valid-type]
101+
def __getattr__(self, attr: str, /) -> Any:
159102
def func(df: DataFrame, *args: Any, **kwargs: Any) -> Any:
160103
parsed_args = parse_args(df, *args)
161104
parsed_kwargs = parse_kwargs(df, **kwargs)
162105
return getattr(self(df), attr)(*parsed_args, **parsed_kwargs)
163106

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

109+
# Namespaces
110+
@property
111+
def dt(self) -> NamespaceExpr:
112+
return NamespaceExpr(self, "dt")
113+
114+
@property
115+
def str(self) -> NamespaceExpr:
116+
return NamespaceExpr(self, "str")
117+
118+
@property
119+
def cat(self) -> NamespaceExpr:
120+
return NamespaceExpr(self, "cat")
121+
122+
@property
123+
def list(self) -> NamespaceExpr:
124+
return NamespaceExpr(self, "list")
125+
126+
@property
127+
def sparse(self) -> NamespaceExpr:
128+
return NamespaceExpr(self, "sparse")
129+
130+
@property
131+
def struct(self) -> NamespaceExpr:
132+
return NamespaceExpr(self, "struct")
133+
166134

167135
class NamespaceExpr:
168136
def __init__(self, func: Callable[[DataFrame], Any], namespace: str) -> None:

0 commit comments

Comments
 (0)