reduce redirection in ops

Author: jbrockmendel

pandas/core/ops.py

@@ -188,6 +188,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
 
@@ -408,75 +480,56 @@ 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']
 
     # 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()}
+    new_methods = {v.__name__: v for k, v in new_methods.items()}
     return new_methods
 
 
@@ -503,12 +556,12 @@ 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)
@@ -552,22 +605,21 @@ 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.
 
     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'],
@@ -626,11 +678,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
@@ -763,11 +817,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):
@@ -890,7 +945,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.
@@ -961,7 +1016,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)
@@ -1042,7 +1098,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)
@@ -1119,7 +1177,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):
@@ -1167,7 +1227,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):  # Another DataFrame
@@ -1200,8 +1263,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 '
@@ -1214,7 +1279,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
 
@@ -1260,7 +1328,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)
 
@@ -1298,18 +1368,19 @@ def f(self, other, axis=0):
                            comp_method=_comp_method_PANEL,
                            bool_method=_arith_method_PANEL)
 
+panel_flex_funcs = dict(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):
@@ -1347,11 +1418,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 (
@@ -1379,3 +1451,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)