Hide keyboard shortcuts

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 

5 

6from pandas.core.dtypes.generic import ABCDataFrame, ABCSeries, ABCSparseArray 

7 

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) 

21 

22 

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. 

27 

28 Parameters 

29 ---------- 

30 cls : class 

31 

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 

39 

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 ) 

55 

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 

70 

71 

72def add_special_arithmetic_methods(cls): 

73 """ 

74 Adds the full suite of special arithmetic methods (``__add__``, 

75 ``__sub__``, etc.) to the class. 

76 

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 

88 

89 def _wrap_inplace_method(method): 

90 """ 

91 return an inplace wrapper for this method 

92 """ 

93 

94 def f(self, other): 

95 result = method(self, other) 

96 

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 ) 

102 

103 return self 

104 

105 name = method.__name__.strip("__") 

106 f.__name__ = f"__i{name}__" 

107 return f 

108 

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 ) 

120 

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 ) 

128 

129 _add_methods(cls, new_methods=new_methods) 

130 

131 

132def add_flex_arithmetic_methods(cls): 

133 """ 

134 Adds the full suite of flex arithmetic methods (``pow``, ``mul``, ``add``) 

135 to the class. 

136 

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_")) 

155 

156 _add_methods(cls, new_methods=new_methods) 

157 

158 

159def _create_methods(cls, arith_method, comp_method, bool_method, special): 

160 # creates actual methods based upon arithmetic, comp and bool method 

161 # constructors. 

162 

163 have_divmod = issubclass(cls, ABCSeries) 

164 # divmod is available for Series 

165 

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) 

192 

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 ) 

203 

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 ) 

216 

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 

223 

224 

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)