diff --git a/pandas/core/ops.py b/pandas/core/ops.py index dff2b6844af94..da65f1f31ed2a 100644 --- a/pandas/core/ops.py +++ b/pandas/core/ops.py @@ -208,6 +208,78 @@ def _get_frame_op_default_axis(name): return 'columns' +def _get_opstr(op, cls): + """ + Find the operation string, if any, to pass to numexpr for this + operation. + + Parameters + ---------- + op : binary operator + cls : class + + Returns + ------- + op_str : string or None + """ + # numexpr is available for non-sparse classes + subtyp = getattr(cls, '_subtyp', '') + use_numexpr = 'sparse' not in subtyp + + if not use_numexpr: + # if we're not using numexpr, then don't pass a str_rep + return None + + return {operator.add: '+', + radd: '+', + operator.mul: '*', + rmul: '*', + operator.sub: '-', + rsub: '-', + operator.truediv: '/', + rtruediv: '/', + operator.floordiv: '//', + rfloordiv: '//', + operator.mod: None, # TODO: Why None for mod but '%' for rmod? + rmod: '%', + operator.pow: '**', + rpow: '**', + operator.eq: '==', + operator.ne: '!=', + operator.le: '<=', + operator.lt: '<', + operator.ge: '>=', + operator.gt: '>', + operator.and_: '&', + rand_: '&', + operator.or_: '|', + ror_: '|', + operator.xor: '^', + rxor: '^', + divmod: None, + rdivmod: None}[op] + + +def _get_op_name(op, special): + """ + Find the name to attach to this method according to conventions + for special and non-special methods. + + Parameters + ---------- + op : binary operator + special : bool + + Returns + ------- + op_name : str + """ + opname = op.__name__.strip('_') + if special: + opname = '__{opname}__'.format(opname=opname) + return opname + + # ----------------------------------------------------------------------------- # Docstring Generation and Templates @@ -501,48 +573,29 @@ def _create_methods(cls, arith_method, comp_method, bool_method, # creates actual methods based upon arithmetic, comp and bool method # constructors. - # numexpr is available for non-sparse classes - subtyp = getattr(cls, '_subtyp', '') - use_numexpr = 'sparse' not in subtyp - have_divmod = issubclass(cls, ABCSeries) # divmod is available for Series and SparseSeries - # if we're not using numexpr, then don't pass a str_rep - if use_numexpr: - op = lambda x: x - else: - op = lambda x: None - if special: - - def names(x): - if x[-1] == "_": - return "__{name}_".format(name=x) - else: - return "__{name}__".format(name=x) - else: - names = lambda x: x - # yapf: disable new_methods = dict( - add=arith_method(operator.add, names('add'), op('+')), - radd=arith_method(radd, names('radd'), op('+')), - sub=arith_method(operator.sub, names('sub'), op('-')), - mul=arith_method(operator.mul, names('mul'), op('*')), - truediv=arith_method(operator.truediv, names('truediv'), op('/')), - floordiv=arith_method(operator.floordiv, names('floordiv'), op('//')), + add=arith_method(cls, operator.add, special), + radd=arith_method(cls, radd, special), + sub=arith_method(cls, operator.sub, special), + mul=arith_method(cls, operator.mul, special), + truediv=arith_method(cls, operator.truediv, special), + floordiv=arith_method(cls, operator.floordiv, special), # Causes a floating point exception in the tests when numexpr enabled, # so for now no speedup - mod=arith_method(operator.mod, names('mod'), None), - pow=arith_method(operator.pow, names('pow'), op('**')), + mod=arith_method(cls, operator.mod, special), + pow=arith_method(cls, operator.pow, special), # not entirely sure why this is necessary, but previously was included # so it's here to maintain compatibility - rmul=arith_method(operator.mul, names('rmul'), op('*')), - rsub=arith_method(rsub, names('rsub'), op('-')), - rtruediv=arith_method(rtruediv, names('rtruediv'), op('/')), - rfloordiv=arith_method(rfloordiv, names('rfloordiv'), op('//')), - rpow=arith_method(rpow, names('rpow'), op('**')), - rmod=arith_method(rmod, names('rmod'), op('%'))) + rmul=arith_method(cls, rmul, special), + rsub=arith_method(cls, rsub, special), + rtruediv=arith_method(cls, rtruediv, special), + rfloordiv=arith_method(cls, rfloordiv, special), + rpow=arith_method(cls, rpow, special), + rmod=arith_method(cls, rmod, special)) # yapf: enable new_methods['div'] = new_methods['truediv'] new_methods['rdiv'] = new_methods['rtruediv'] @@ -550,26 +603,30 @@ def names(x): # Comp methods never had a default axis set if comp_method: new_methods.update(dict( - eq=comp_method(operator.eq, names('eq'), op('==')), - ne=comp_method(operator.ne, names('ne'), op('!=')), - lt=comp_method(operator.lt, names('lt'), op('<')), - gt=comp_method(operator.gt, names('gt'), op('>')), - le=comp_method(operator.le, names('le'), op('<=')), - ge=comp_method(operator.ge, names('ge'), op('>=')))) + eq=comp_method(cls, operator.eq, special), + ne=comp_method(cls, operator.ne, special), + lt=comp_method(cls, operator.lt, special), + gt=comp_method(cls, operator.gt, special), + le=comp_method(cls, operator.le, special), + ge=comp_method(cls, operator.ge, special))) if bool_method: new_methods.update( - dict(and_=bool_method(operator.and_, names('and_'), op('&')), - or_=bool_method(operator.or_, names('or_'), op('|')), + dict(and_=bool_method(cls, operator.and_, special), + or_=bool_method(cls, operator.or_, special), # For some reason ``^`` wasn't used in original. - xor=bool_method(operator.xor, names('xor'), op('^')), - rand_=bool_method(rand_, names('rand_'), op('&')), - ror_=bool_method(ror_, names('ror_'), op('|')), - rxor=bool_method(rxor, names('rxor'), op('^')))) + xor=bool_method(cls, operator.xor, special), + rand_=bool_method(cls, rand_, special), + ror_=bool_method(cls, ror_, special), + rxor=bool_method(cls, rxor, special))) if have_divmod: # divmod doesn't have an op that is supported by numexpr - new_methods['divmod'] = arith_method(divmod, names('divmod'), None) + new_methods['divmod'] = arith_method(cls, divmod, special) - new_methods = {names(k): v for k, v in new_methods.items()} + if special: + dunderize = lambda x: '__{name}__'.format(name=x.strip('_')) + else: + dunderize = lambda x: x + new_methods = {dunderize(k): v for k, v in new_methods.items()} return new_methods @@ -596,16 +653,15 @@ def add_special_arithmetic_methods(cls, arith_method=None, Parameters ---------- arith_method : function (optional) - factory for special arithmetic methods, with op string: - f(op, name, str_rep) + factory for special arithmetic methods: + f(cls, op, special) comp_method : function (optional) - factory for rich comparison - signature: f(op, name, str_rep) + factory for rich comparison - signature: f(cls, op, special) bool_method : function (optional) - factory for boolean methods - signature: f(op, name, str_rep) + factory for boolean methods - signature: f(cls, op, special) """ new_methods = _create_methods(cls, arith_method, comp_method, bool_method, special=True) - # inplace operators (I feel like these should get passed an `inplace=True` # or just be removed @@ -645,8 +701,7 @@ def f(self, other): add_methods(cls, new_methods=new_methods) -def add_flex_arithmetic_methods(cls, flex_arith_method, - flex_comp_method=None, flex_bool_method=None): +def add_flex_arithmetic_methods(cls, flex_arith_method, flex_comp_method=None): """ Adds the full suite of flex arithmetic methods (``pow``, ``mul``, ``add``) to the class. @@ -654,13 +709,13 @@ def add_flex_arithmetic_methods(cls, flex_arith_method, Parameters ---------- flex_arith_method : function - factory for flex arithmetic methods, with op string: - f(op, name, str_rep) + factory for flex arithmetic methods: + f(cls, op, special) flex_comp_method : function, optional, - factory for rich comparison - signature: f(op, name, str_rep) + factory for rich comparison - signature: f(cls, op, special) """ new_methods = _create_methods(cls, flex_arith_method, - flex_comp_method, flex_bool_method, + flex_comp_method, bool_method=None, special=False) new_methods.update(dict(multiply=new_methods['mul'], subtract=new_methods['sub'], @@ -719,11 +774,13 @@ def _construct_divmod_result(left, result, index, name, dtype): ) -def _arith_method_SERIES(op, name, str_rep): +def _arith_method_SERIES(cls, op, special): """ Wrapper function for Series arithmetic operations, to avoid code duplication. """ + str_rep = _get_opstr(op, cls) + name = _get_op_name(op, special) eval_kwargs = _gen_eval_kwargs(name) fill_zeros = _gen_fill_zeros(name) construct_result = (_construct_divmod_result @@ -856,11 +913,12 @@ def _comp_method_OBJECT_ARRAY(op, x, y): return result -def _comp_method_SERIES(op, name, str_rep): +def _comp_method_SERIES(cls, op, special): """ Wrapper function for Series arithmetic operations, to avoid code duplication. """ + name = _get_op_name(op, special) masker = _gen_eval_kwargs(name).get('masker', False) def na_op(x, y): @@ -995,7 +1053,7 @@ def wrapper(self, other, axis=None): return wrapper -def _bool_method_SERIES(op, name, str_rep): +def _bool_method_SERIES(cls, op, special): """ Wrapper function for Series arithmetic operations, to avoid code duplication. @@ -1066,7 +1124,8 @@ def wrapper(self, other): return wrapper -def _flex_method_SERIES(op, name, str_rep): +def _flex_method_SERIES(cls, op, special): + name = _get_op_name(op, special) doc = _make_flex_doc(name, 'series') @Appender(doc) @@ -1192,7 +1251,9 @@ def to_series(right): return right -def _arith_method_FRAME(op, name, str_rep=None): +def _arith_method_FRAME(cls, op, special): + str_rep = _get_opstr(op, cls) + name = _get_op_name(op, special) eval_kwargs = _gen_eval_kwargs(name) fill_zeros = _gen_fill_zeros(name) default_axis = _get_frame_op_default_axis(name) @@ -1270,7 +1331,9 @@ def f(self, other, axis=default_axis, level=None, fill_value=None): return f -def _flex_comp_method_FRAME(op, name, str_rep=None): +def _flex_comp_method_FRAME(cls, op, special): + str_rep = _get_opstr(op, cls) + name = _get_op_name(op, special) default_axis = _get_frame_op_default_axis(name) def na_op(x, y): @@ -1306,7 +1369,10 @@ def f(self, other, axis=default_axis, level=None): return f -def _comp_method_FRAME(func, name, str_rep): +def _comp_method_FRAME(cls, func, special): + str_rep = _get_opstr(func, cls) + name = _get_op_name(func, special) + @Appender('Wrapper for comparison method {name}'.format(name=name)) def f(self, other): if isinstance(other, ABCDataFrame): @@ -1345,8 +1411,10 @@ def f(self, other): # ----------------------------------------------------------------------------- # Panel -def _arith_method_PANEL(op, name, str_rep=None): +def _arith_method_PANEL(cls, op, special): # work only for scalars + name = _get_op_name(op, special) + def f(self, other): if not is_scalar(other): raise ValueError('Simple arithmetic with {name} can only be ' @@ -1359,7 +1427,10 @@ def f(self, other): return f -def _comp_method_PANEL(op, name, str_rep=None): +def _comp_method_PANEL(cls, op, special): + str_rep = _get_opstr(op, cls) + name = _get_op_name(op, special) + def na_op(x, y): import pandas.core.computation.expressions as expressions @@ -1389,7 +1460,9 @@ def f(self, other, axis=None): return f -def _flex_method_PANEL(op, name, str_rep=None): +def _flex_method_PANEL(cls, op, special): + str_rep = _get_opstr(op, cls) + name = _get_op_name(op, special) eval_kwargs = _gen_eval_kwargs(name) fill_zeros = _gen_fill_zeros(name) @@ -1427,18 +1500,19 @@ def f(self, other, axis=0): comp_method=_comp_method_PANEL, bool_method=_arith_method_PANEL) +panel_flex_funcs = dict(flex_arith_method=_flex_method_PANEL, + flex_comp_method=_comp_method_PANEL) # ----------------------------------------------------------------------------- # Sparse -def _arith_method_SPARSE_SERIES(op, name, str_rep=None): +def _arith_method_SPARSE_SERIES(cls, op, special): """ Wrapper function for Series arithmetic operations, to avoid code duplication. - - str_rep is not used, but is present for compatibility. """ + name = _get_op_name(op, special) def wrapper(self, other): if isinstance(other, ABCDataFrame): @@ -1476,11 +1550,12 @@ def _sparse_series_op(left, right, op, name): return left._constructor(result, index=new_index, name=new_name) -def _arith_method_SPARSE_ARRAY(op, name, str_rep=None): +def _arith_method_SPARSE_ARRAY(cls, op, special): """ Wrapper function for Series arithmetic operations, to avoid code duplication. """ + name = _get_op_name(op, special) def wrapper(self, other): from pandas.core.sparse.array import ( @@ -1508,3 +1583,12 @@ def wrapper(self, other): name = name[2:-2] wrapper.__name__ = name return wrapper + + +sparse_array_special_funcs = dict(arith_method=_arith_method_SPARSE_ARRAY, + comp_method=_arith_method_SPARSE_ARRAY, + bool_method=_arith_method_SPARSE_ARRAY) + +sparse_series_special_funcs = dict(arith_method=_arith_method_SPARSE_SERIES, + comp_method=_arith_method_SPARSE_SERIES, + bool_method=None) diff --git a/pandas/core/panel.py b/pandas/core/panel.py index 7f973992fb07f..3be1e3ef8734d 100644 --- a/pandas/core/panel.py +++ b/pandas/core/panel.py @@ -1528,8 +1528,7 @@ def _extract_axis(self, data, axis=0, intersect=False): 'minor_axis': 'columns'}) ops.add_special_arithmetic_methods(Panel, **ops.panel_special_funcs) -ops.add_flex_arithmetic_methods(Panel, ops._flex_method_PANEL, - flex_comp_method=ops._comp_method_PANEL) +ops.add_flex_arithmetic_methods(Panel, **ops.panel_flex_funcs) Panel._add_numeric_operations() diff --git a/pandas/core/sparse/array.py b/pandas/core/sparse/array.py index 3cbae717d0e07..4f7152666f7bf 100644 --- a/pandas/core/sparse/array.py +++ b/pandas/core/sparse/array.py @@ -844,6 +844,4 @@ def _make_index(length, indices, kind): ops.add_special_arithmetic_methods(SparseArray, - arith_method=ops._arith_method_SPARSE_ARRAY, - comp_method=ops._arith_method_SPARSE_ARRAY, - bool_method=ops._arith_method_SPARSE_ARRAY) + **ops.sparse_array_special_funcs) diff --git a/pandas/core/sparse/series.py b/pandas/core/sparse/series.py index 62a467bec2683..335a4c80adc63 100644 --- a/pandas/core/sparse/series.py +++ b/pandas/core/sparse/series.py @@ -817,6 +817,4 @@ def from_coo(cls, A, dense_index=False): # overwrite basic arithmetic to use SparseSeries version # force methods to overwrite previous definitions. ops.add_special_arithmetic_methods(SparseSeries, - ops._arith_method_SPARSE_SERIES, - comp_method=ops._arith_method_SPARSE_SERIES, - bool_method=None) + **ops.sparse_series_special_funcs)