Coverage for /home/martinb/.local/share/virtualenvs/camcops/lib/python3.6/site-packages/pandas/core/ops/methods.py : 95%

Hot-keys on this page
r m x p toggle line displays
j k next/prev highlighted chunk
0 (zero) top of page
1 (one) first highlighted chunk
1"""
2Functions to generate methods and pin them to the appropriate classes.
3"""
4import operator
6from pandas.core.dtypes.generic import ABCDataFrame, ABCSeries, ABCSparseArray
8from pandas.core.ops.roperator import (
9 radd,
10 rand_,
11 rdivmod,
12 rfloordiv,
13 rmod,
14 rmul,
15 ror_,
16 rpow,
17 rsub,
18 rtruediv,
19 rxor,
20)
23def _get_method_wrappers(cls):
24 """
25 Find the appropriate operation-wrappers to use when defining flex/special
26 arithmetic, boolean, and comparison operations with the given class.
28 Parameters
29 ----------
30 cls : class
32 Returns
33 -------
34 arith_flex : function or None
35 comp_flex : function or None
36 arith_special : function
37 comp_special : function
38 bool_special : function
40 Notes
41 -----
42 None is only returned for SparseArray
43 """
44 # TODO: make these non-runtime imports once the relevant functions
45 # are no longer in __init__
46 from pandas.core.ops import (
47 _arith_method_FRAME,
48 _arith_method_SERIES,
49 _bool_method_SERIES,
50 _comp_method_FRAME,
51 _comp_method_SERIES,
52 _flex_comp_method_FRAME,
53 _flex_method_SERIES,
54 )
56 if issubclass(cls, ABCSeries):
57 # Just Series
58 arith_flex = _flex_method_SERIES
59 comp_flex = _flex_method_SERIES
60 arith_special = _arith_method_SERIES
61 comp_special = _comp_method_SERIES
62 bool_special = _bool_method_SERIES
63 elif issubclass(cls, ABCDataFrame):
64 arith_flex = _arith_method_FRAME
65 comp_flex = _flex_comp_method_FRAME
66 arith_special = _arith_method_FRAME
67 comp_special = _comp_method_FRAME
68 bool_special = _arith_method_FRAME
69 return arith_flex, comp_flex, arith_special, comp_special, bool_special
72def add_special_arithmetic_methods(cls):
73 """
74 Adds the full suite of special arithmetic methods (``__add__``,
75 ``__sub__``, etc.) to the class.
77 Parameters
78 ----------
79 cls : class
80 special methods will be defined and pinned to this class
81 """
82 _, _, arith_method, comp_method, bool_method = _get_method_wrappers(cls)
83 new_methods = _create_methods(
84 cls, arith_method, comp_method, bool_method, special=True
85 )
86 # inplace operators (I feel like these should get passed an `inplace=True`
87 # or just be removed
89 def _wrap_inplace_method(method):
90 """
91 return an inplace wrapper for this method
92 """
94 def f(self, other):
95 result = method(self, other)
97 # this makes sure that we are aligned like the input
98 # we are updating inplace so we want to ignore is_copy
99 self._update_inplace(
100 result.reindex_like(self, copy=False)._data, verify_is_copy=False
101 )
103 return self
105 name = method.__name__.strip("__")
106 f.__name__ = f"__i{name}__"
107 return f
109 new_methods.update(
110 dict(
111 __iadd__=_wrap_inplace_method(new_methods["__add__"]),
112 __isub__=_wrap_inplace_method(new_methods["__sub__"]),
113 __imul__=_wrap_inplace_method(new_methods["__mul__"]),
114 __itruediv__=_wrap_inplace_method(new_methods["__truediv__"]),
115 __ifloordiv__=_wrap_inplace_method(new_methods["__floordiv__"]),
116 __imod__=_wrap_inplace_method(new_methods["__mod__"]),
117 __ipow__=_wrap_inplace_method(new_methods["__pow__"]),
118 )
119 )
121 new_methods.update(
122 dict(
123 __iand__=_wrap_inplace_method(new_methods["__and__"]),
124 __ior__=_wrap_inplace_method(new_methods["__or__"]),
125 __ixor__=_wrap_inplace_method(new_methods["__xor__"]),
126 )
127 )
129 _add_methods(cls, new_methods=new_methods)
132def add_flex_arithmetic_methods(cls):
133 """
134 Adds the full suite of flex arithmetic methods (``pow``, ``mul``, ``add``)
135 to the class.
137 Parameters
138 ----------
139 cls : class
140 flex methods will be defined and pinned to this class
141 """
142 flex_arith_method, flex_comp_method, _, _, _ = _get_method_wrappers(cls)
143 new_methods = _create_methods(
144 cls, flex_arith_method, flex_comp_method, bool_method=None, special=False
145 )
146 new_methods.update(
147 dict(
148 multiply=new_methods["mul"],
149 subtract=new_methods["sub"],
150 divide=new_methods["div"],
151 )
152 )
153 # opt out of bool flex methods for now
154 assert not any(kname in new_methods for kname in ("ror_", "rxor", "rand_"))
156 _add_methods(cls, new_methods=new_methods)
159def _create_methods(cls, arith_method, comp_method, bool_method, special):
160 # creates actual methods based upon arithmetic, comp and bool method
161 # constructors.
163 have_divmod = issubclass(cls, ABCSeries)
164 # divmod is available for Series
166 new_methods = dict(
167 add=arith_method(cls, operator.add, special),
168 radd=arith_method(cls, radd, special),
169 sub=arith_method(cls, operator.sub, special),
170 mul=arith_method(cls, operator.mul, special),
171 truediv=arith_method(cls, operator.truediv, special),
172 floordiv=arith_method(cls, operator.floordiv, special),
173 # Causes a floating point exception in the tests when numexpr enabled,
174 # so for now no speedup
175 mod=arith_method(cls, operator.mod, special),
176 pow=arith_method(cls, operator.pow, special),
177 # not entirely sure why this is necessary, but previously was included
178 # so it's here to maintain compatibility
179 rmul=arith_method(cls, rmul, special),
180 rsub=arith_method(cls, rsub, special),
181 rtruediv=arith_method(cls, rtruediv, special),
182 rfloordiv=arith_method(cls, rfloordiv, special),
183 rpow=arith_method(cls, rpow, special),
184 rmod=arith_method(cls, rmod, special),
185 )
186 new_methods["div"] = new_methods["truediv"]
187 new_methods["rdiv"] = new_methods["rtruediv"]
188 if have_divmod:
189 # divmod doesn't have an op that is supported by numexpr
190 new_methods["divmod"] = arith_method(cls, divmod, special)
191 new_methods["rdivmod"] = arith_method(cls, rdivmod, special)
193 new_methods.update(
194 dict(
195 eq=comp_method(cls, operator.eq, special),
196 ne=comp_method(cls, operator.ne, special),
197 lt=comp_method(cls, operator.lt, special),
198 gt=comp_method(cls, operator.gt, special),
199 le=comp_method(cls, operator.le, special),
200 ge=comp_method(cls, operator.ge, special),
201 )
202 )
204 if bool_method:
205 new_methods.update(
206 dict(
207 and_=bool_method(cls, operator.and_, special),
208 or_=bool_method(cls, operator.or_, special),
209 # For some reason ``^`` wasn't used in original.
210 xor=bool_method(cls, operator.xor, special),
211 rand_=bool_method(cls, rand_, special),
212 ror_=bool_method(cls, ror_, special),
213 rxor=bool_method(cls, rxor, special),
214 )
215 )
217 if special:
218 dunderize = lambda x: f"__{x.strip('_')}__"
219 else:
220 dunderize = lambda x: x
221 new_methods = {dunderize(k): v for k, v in new_methods.items()}
222 return new_methods
225def _add_methods(cls, new_methods):
226 for name, method in new_methods.items():
227 # For most methods, if we find that the class already has a method
228 # of the same name, it is OK to over-write it. The exception is
229 # inplace methods (__iadd__, __isub__, ...) for SparseArray, which
230 # retain the np.ndarray versions.
231 force = not (issubclass(cls, ABCSparseArray) and name.startswith("__i"))
232 if force or name not in cls.__dict__:
233 setattr(cls, name, method)