Add CompWrapper to clean code

Author: makbigc

pandas/core/arrays/__init__.py

@@ -1,7 +1,8 @@
 from .array_ import array  # noqa
 from .base import (ExtensionArray,    # noqa
                    ExtensionOpsMixin,
-                   ExtensionScalarOpsMixin)
+                   ExtensionScalarOpsMixin,
+                   CompWrapper)
 from .categorical import Categorical  # noqa
 from .datetimes import DatetimeArray  # noqa
 from .interval import IntervalArray  # noqa

pandas/core/arrays/base.py

@@ -6,6 +6,7 @@
    without warning.
 """
 import operator
+from functools import wraps
 
 import numpy as np
 
@@ -15,7 +16,7 @@
 from pandas.util._decorators import Appender, Substitution
 
 from pandas.core.dtypes.common import is_list_like
-from pandas.core.dtypes.generic import ABCIndexClass, ABCSeries
+from pandas.core.dtypes.generic import ABCDataFrame, ABCIndexClass, ABCSeries
 from pandas.core.dtypes.missing import isna
 
 from pandas.core import ops
@@ -1118,3 +1119,83 @@ def _create_arithmetic_method(cls, op):
     @classmethod
     def _create_comparison_method(cls, op):
         return cls._create_method(op, coerce_to_dtype=False)
+
+'''
+def validate_comp_other(comp, list_to_array=False, validate_len=False,
+                        zerodim=False, inst_from_senior_cls=False):
+    def wrapper(self, other):
+        if list_to_array is True:
+            if is_list_like(other):
+                other = np.asarray(other)
+
+        if validate_len is True:
+            if is_list_like(other) and len(other) != len(self):
+                raise ValueError("Lenghts must match")
+
+        if zerodim is True:
+            import pandas._libs as lib
+            other = lib.item_from_zerodim(other)
+
+        if inst_from_senior_cls is True:
+            if isinstance(other, (ABCDataFrame, ABCSeries, ABCIndexClass)):
+                # Rely on pandas to unbox and dispatch to us.
+                return NotImplemented
+
+        comp(self, other)
+    return wrapper
+'''
+
+
+class CompWrapper(object):
+    __key__ = ['list_to_array', 'validate_len',
+               'zerodim', 'inst_from_senior_cls']
+
+    def __init__(self,
+                 list_to_array=None,
+                 validate_len=None,
+                 zerodim=None,
+                 inst_from_senior_cls=None):
+        self.list_to_array = list_to_array
+        self.validate_len = validate_len
+        self.zerodim = zerodim
+        self.inst_from_senior_cls = inst_from_senior_cls
+
+    def _list_to_array(self, comp):
+        @wraps(comp)
+        def wrapper(comp_self, comp_other):
+            if is_list_like(comp_other):
+                comp_other = np.asarray(comp_other)
+            return comp(comp_self, comp_other)
+        return wrapper
+
+    def _validate_len(self, comp):
+        @wraps(comp)
+        def wrapper(comp_self, comp_other):
+            if is_list_like(comp_other) and len(comp_other) != len(comp_self):
+                raise ValueError("Lengths must match to compare")
+            return comp(comp_self, comp_other)
+        return wrapper
+
+    def _zerodim(self, comp):
+        @wraps(comp)
+        def wrapper(comp_self, comp_other):
+            from pandas._libs import lib
+            comp_other = lib.item_from_zerodim(comp_other)
+            return comp(comp_self, comp_other)
+        return wrapper
+
+    def _inst_from_senior_cls(self, comp):
+        @wraps(comp)
+        def wrapper(comp_self, comp_other):
+            if isinstance(comp_other, (ABCDataFrame,
+                                       ABCSeries, ABCIndexClass)):
+                # Rely on pandas to unbox and dispatch to us.
+                return NotImplemented
+            return comp(comp_self, comp_other)
+        return wrapper
+
+    def __call__(self, comp):
+        for key in CompWrapper.__key__:
+                if getattr(self, key) is True:
+                    comp = getattr(self, '_' + key)(comp)
+        return comp

pandas/core/arrays/datetimes.py

@@ -20,12 +20,12 @@
     is_string_dtype, is_timedelta64_dtype, pandas_dtype)
 from pandas.core.dtypes.dtypes import DatetimeTZDtype
 from pandas.core.dtypes.generic import (
-    ABCDataFrame, ABCIndexClass, ABCPandasArray, ABCSeries)
+    ABCIndexClass, ABCPandasArray, ABCSeries)
 from pandas.core.dtypes.missing import isna
 
 from pandas.core import ops
 from pandas.core.algorithms import checked_add_with_arr
-from pandas.core.arrays import datetimelike as dtl
+from pandas.core.arrays import datetimelike as dtl, CompWrapper
 from pandas.core.arrays._ranges import generate_regular_range
 import pandas.core.common as com
 
@@ -129,12 +129,8 @@ def _dt_array_cmp(cls, op):
     opname = '__{name}__'.format(name=op.__name__)
     nat_result = True if opname == '__ne__' else False
 
+    @CompWrapper(inst_from_senior_cls=True, validate_len=True, zerodim=True)
     def wrapper(self, other):
-        if isinstance(other, (ABCDataFrame, ABCSeries, ABCIndexClass)):
-            return NotImplemented
-
-        other = lib.item_from_zerodim(other)
-
         if isinstance(other, (datetime, np.datetime64, compat.string_types)):
             if isinstance(other, (datetime, np.datetime64)):
                 # GH#18435 strings get a pass from tzawareness compat
@@ -151,8 +147,8 @@ def wrapper(self, other):
                 result.fill(nat_result)
         elif lib.is_scalar(other) or np.ndim(other) == 0:
             return ops.invalid_comparison(self, other, op)
-        elif len(other) != len(self):
-            raise ValueError("Lengths must match")
+        #elif len(other) != len(self):
+        #    raise ValueError("Lengths must match")
         else:
             if isinstance(other, list):
                 try:

pandas/core/arrays/integer.py

@@ -18,7 +18,7 @@
 from pandas.core.dtypes.missing import isna, notna
 
 from pandas.core import nanops
-from pandas.core.arrays import ExtensionArray, ExtensionOpsMixin
+from pandas.core.arrays import ExtensionArray, ExtensionOpsMixin, CompWrapper
 from pandas.core.tools.numeric import to_numeric
 
 
@@ -496,22 +496,17 @@ def _values_for_argsort(self):
 
     @classmethod
     def _create_comparison_method(cls, op):
+        @CompWrapper(validate_len=True, inst_from_senior_cls=True)
         def cmp_method(self, other):
 
             op_name = op.__name__
             mask = None
 
-            if isinstance(other, (ABCSeries, ABCIndexClass)):
-                # Rely on pandas to unbox and dispatch to us.
-                return NotImplemented
-
             if isinstance(other, IntegerArray):
                 other, mask = other._data, other._mask
 
             elif is_list_like(other):
                 other = np.asarray(other)
-                if other.ndim > 0 and len(self) != len(other):
-                    raise ValueError('Lengths must match to compare')
 
             other = lib.item_from_zerodim(other)
 

pandas/core/arrays/period.py

@@ -15,14 +15,15 @@
 
 from pandas.core.dtypes.common import (
     _TD_DTYPE, ensure_object, is_datetime64_dtype, is_float_dtype,
-    is_list_like, is_period_dtype, pandas_dtype)
+    is_period_dtype, pandas_dtype)
 from pandas.core.dtypes.dtypes import PeriodDtype
 from pandas.core.dtypes.generic import (
-    ABCDataFrame, ABCIndexClass, ABCPeriodIndex, ABCSeries)
+    ABCIndexClass, ABCPeriodIndex, ABCSeries)
 from pandas.core.dtypes.missing import isna, notna
 
 import pandas.core.algorithms as algos
 from pandas.core.arrays import datetimelike as dtl
+from pandas.core.arrays.base import CompWrapper
 import pandas.core.common as com
 
 from pandas.tseries import frequencies
@@ -47,15 +48,10 @@ def _period_array_cmp(cls, op):
     opname = '__{name}__'.format(name=op.__name__)
     nat_result = True if opname == '__ne__' else False
 
+    @CompWrapper(validate_len=True, inst_from_senior_cls=True)
     def wrapper(self, other):
         op = getattr(self.asi8, opname)
 
-        if isinstance(other, (ABCDataFrame, ABCSeries, ABCIndexClass)):
-            return NotImplemented
-
-        if is_list_like(other) and len(other) != len(self):
-            raise ValueError("Lengths must match")
-
         if isinstance(other, Period):
             self._check_compatible_with(other)
 

pandas/core/arrays/timedeltas.py

@@ -14,6 +14,7 @@
 import pandas.compat as compat
 from pandas.util._decorators import Appender
 
+from pandas.core.arrays import CompWrapper
 from pandas.core.dtypes.common import (
     _NS_DTYPE, _TD_DTYPE, ensure_int64, is_datetime64_dtype, is_float_dtype,
     is_integer_dtype, is_list_like, is_object_dtype, is_scalar,
@@ -63,10 +64,8 @@ def _td_array_cmp(cls, op):
     opname = '__{name}__'.format(name=op.__name__)
     nat_result = True if opname == '__ne__' else False
 
+    @CompWrapper(validate_len=True, inst_from_senior_cls=True)
     def wrapper(self, other):
-        if isinstance(other, (ABCDataFrame, ABCSeries, ABCIndexClass)):
-            return NotImplemented
-
         if _is_convertible_to_td(other) or other is NaT:
             try:
                 other = Timedelta(other)
@@ -81,9 +80,6 @@ def wrapper(self, other):
         elif not is_list_like(other):
             return ops.invalid_comparison(self, other, op)
 
-        elif len(other) != len(self):
-            raise ValueError("Lengths must match")
-
         else:
             try:
                 other = type(self)._from_sequence(other)._data