REF: Split argsort into two parts

Author: TomAugspurger

pandas/core/arrays/base.py

@@ -2,6 +2,7 @@
 import numpy as np
 
 from pandas.errors import AbstractMethodError
+from pandas.compat.numpy import function as nv
 
 _not_implemented_message = "{} does not implement {}."
 
@@ -216,11 +217,24 @@ def isna(self):
         """
         raise AbstractMethodError(self)
 
-    def argsort(self, axis=-1, kind='quicksort', order=None):
+    def _values_for_argsort(self):
+        # type: () -> ndarray
+        """Get the ndarray to be passed to np.argsort.
+
+        This is called from within 'ExtensionArray.argsort'.
+
+        Returns
+        -------
+        values : ndarray
+        """
+        return np.array(self)
+
+    def argsort(self, ascending=True, kind='quicksort', *args, **kwargs):
         """Returns the indices that would sort this array.
 
         Parameters
         ----------
+        ascending : bool, default True
         axis : int or None, optional
             Axis along which to sort. ExtensionArrays are 1-dimensional,
             so this is only included for compatibility with NumPy.
@@ -233,9 +247,18 @@ def argsort(self, axis=-1, kind='quicksort', order=None):
         -------
         index_array : ndarray
             Array of indices that sort ``self``.
-
         """
-        return np.array(self).argsort(kind=kind)
+        # Implementor note: You have two places to override the behavior of
+        # argsort.
+        # 1. _values_for_argsort : construct the values passed to np.argsort
+        # 2. argsort : total control over sorting.
+
+        ascending = nv.validate_argsort_with_ascending(ascending, args, kwargs)
+        values = self._values_for_argsort()
+        result = np.argsort(values, kind=kind, **kwargs)
+        if not ascending:
+            result = result[::-1]
+        return result
 
     # ------------------------------------------------------------------------
     # Indexing methods

pandas/core/arrays/categorical.py

@@ -1386,6 +1386,9 @@ def check_for_ordered(self, op):
                             "you can use .as_ordered() to change the "
                             "Categorical to an ordered one\n".format(op=op))
 
+    def _values_for_argsort(self):
+        return self._codes.copy()
+
     def argsort(self, ascending=True, kind='quicksort', *args, **kwargs):
         """
         Returns the indices that would sort the Categorical instance if
@@ -1406,11 +1409,9 @@ def argsort(self, ascending=True, kind='quicksort', *args, **kwargs):
         --------
         numpy.ndarray.argsort
         """
-        ascending = nv.validate_argsort_with_ascending(ascending, args, kwargs)
-        result = np.argsort(self._codes.copy(), kind=kind, **kwargs)
-        if not ascending:
-            result = result[::-1]
-        return result
+        # Keep the implementation here just for the docstring.
+        return super(Categorical, self).argsort(ascending=ascending, kind=kind,
+                                                *args, **kwargs)
 
     def sort_values(self, inplace=False, ascending=True, na_position='last'):
         """ Sorts the Categorical by category value returning a new