temporarily reverted dict keys

Author: Jacob Austin

pandas/core/series.py

@@ -3052,7 +3052,8 @@ def _try_kind_sort(arr):
         ser = ensure_key_mapped(self, key=key)
         sorted_index = np.empty(len(self), dtype=np.int32)
 
-        bad = isna(ser._values)
+        sort_values = ser._values
+        bad = isna(sort_values)
 
         good = ~bad
         idx = ibase.default_index(len(self))

pandas/core/sorting.py

@@ -1,5 +1,5 @@
 """ miscellaneous sorting / groupby utilities """
-from typing import Callable, Dict, Union
+from typing import Callable, Optional
 
 import numpy as np
 
@@ -299,32 +299,20 @@ def nargsort(
     return indexer
 
 
-def apply_key_name(values, key, name):
-    if isinstance(key, dict):
-        key = key.get(name, None)
-
-    if key is None:
-        return values
-
-    return key(values)
-
-
-def ensure_key_mapped_dataframe(df, key, levels=None, axis=0):
+def ensure_key_mapped_dataframe(df, key: Callable, levels=None, axis=0):
     """
     Returns a new DataFrame in which key has been applied
     to all levels specified in level (or all levels if level
-    is None). Used for key sorting for DataFrames.
+    is None) along an axis. Used for key sorting for DataFrames.
 
     Parameters
     ----------
     df : DataFrame
         DataFrame to which to apply the key function on the
         specified levels.
-    key : Callable or Dict[Any, Callable]
-        If Callable, function that takes a Series and returns
-        a Series of the same shape. This key is applied to each
-        level separately. If dict, name or index of each column
-        or row is used to index the key object to get a Callable.
+    key : Callable
+        Function that takes a Series and returns a Series of
+        the same shape. This key is applied to each level separately.
     levels : list-like, int or str, default None
         Level or list of levels to apply the key function to.
         If None, key function is applied to all levels. Other
@@ -345,9 +333,9 @@ def ensure_key_mapped_dataframe(df, key, levels=None, axis=0):
         )
 
     if axis == 0:
-        axis_levels = df.columns._values
+        axis_levels = list(df.columns._values)  # makes mypy happy
     else:
-        axis_levels = df.index._values
+        axis_levels = list(df.index._values)
 
     if levels is not None:
         if isinstance(levels, (str, int)):
@@ -357,17 +345,16 @@ def ensure_key_mapped_dataframe(df, key, levels=None, axis=0):
     else:
         sort_levels = axis_levels
 
-    new_levels = [
-        ensure_key_mapped(
-            Series(df._get_label_or_level_values(name, axis=axis), name=name),
-            key,
-            name=name,
-        )
-        if name in sort_levels
-        else df._get_label_or_level_values(name, axis=axis)
+    values = [
+        (name, Series(df._get_label_or_level_values(name, axis=axis), name=name))
         for name in axis_levels
     ]
 
+    new_levels = [
+        ensure_key_mapped(series, key) if name in sort_levels else series
+        for (name, series) in values
+    ]
+
     if axis == 0:
         new_df = DataFrame._from_arrays(new_levels, df.columns, df.index)
     else:
@@ -419,11 +406,11 @@ def ensure_key_mapped_multiindex(index, key: Callable, levels=None):
     else:
         sort_levels = list(range(index.nlevels))  # satisfies mypy
 
+    values = [(level, index._get_level_values(level)) for level in range(index.nlevels)]
+
     mapped = [
-        ensure_key_mapped(index._get_level_values(level), key, name=level)
-        if level in sort_levels
-        else index._get_level_values(level)
-        for level in range(index.nlevels)
+        ensure_key_mapped(idx, key) if level in sort_levels else idx
+        for (level, idx) in values
     ]
 
     labels = MultiIndex.from_arrays(mapped)
@@ -432,7 +419,7 @@ def ensure_key_mapped_multiindex(index, key: Callable, levels=None):
 
 
 def ensure_key_mapped(
-    values, key: Union[Dict, Callable], levels=None, name=None, axis=0,
+    values, key: Optional[Callable], levels=None, axis=0,
 ):
 
     """
@@ -443,19 +430,19 @@ def ensure_key_mapped(
     Parameters
     ----------
     values : Series, DataFrame, Index subclass, or ndarray
-    key : Union[Callable, Dict[Union[str, int], Callable]
-        key to be called on the values array. If dict, indexed by
-        name or number of columns in values. Dict supported only for
-        DataFrame or MultiIndex. Expected to return an object of the
-        same shape and compatible with the original type.
+    key : Optional[Callable]
+        key to be called on the values array. Expected to
+        take a level or column from values and return an
+        object of the same shape and compatible with the original type.
+        For MultiIndex and DataFrame, applied to rows or columns of the
+        values array. For Series and MultiIndex, applied directly.
     levels :  list-like, int or str, default None
         For MultiIndex values, level or list of levels to apply the key
         function to. If None, key function is applied to all levels. Other
         levels are left unchanged.
-    name : str or int, default None
-        Name used to index the key function if a dictionary.
     axis : int, default 0
-        Axis to use for applying the key to DataFrame values level by level.
+        Axis to use for applying the key to DataFrame values. 0 applies the
+        key to columns, 1 to rows.
     """
     from pandas.core.indexes.api import Index
     from pandas import DataFrame
@@ -469,18 +456,18 @@ def ensure_key_mapped(
     if isinstance(values, DataFrame):  # apply the key to select levels
         return ensure_key_mapped_dataframe(values, key, levels=levels, axis=axis)
 
-    result = apply_key_name(values.copy(), key, name)
+    result = key(values.copy())
 
     if len(result) != len(values):
         raise ValueError(
             "User-provided `key` function must not change the shape of the array."
         )
 
     try:
-        if isinstance(values, Index):
+        if isinstance(values, Index):  # allow a new Index class
             result = Index(result)
         else:
-            result = type(values)(result)
+            result = type(values)(result)  # try to recover otherwise
     except TypeError:
         raise TypeError(
             "User-provided `key` function returned an invalid type {} \

pandas/tests/frame/methods/test_sort_index.py

@@ -396,20 +396,3 @@ def test_changes_length_raises(self):
         df = pd.DataFrame({"A": [1, 2, 3]})
         with pytest.raises(ValueError, match="change the shape"):
             df.sort_index(key=lambda x: x[:1])
-
-    def test_sort_index_key_dict(self):
-        df = DataFrame({0: ["Hello", "goodbye"], 1: [0, 1], 2: [3, 4]}).set_index(
-            [0, 1]
-        )
-
-        result = df.sort_index(level=0, key=lambda col: col.str.lower())
-        expected = df[::-1]
-        tm.assert_frame_equal(result, expected)
-
-        result = df.sort_index(level=[0, 1], key={0: lambda col: col.str.lower()})
-        expected = df[::-1]
-        tm.assert_frame_equal(result, expected)
-
-        result = df.sort_index([0, 1], key={1: lambda col: -col})
-        expected = df
-        tm.assert_frame_equal(result, expected)

pandas/tests/frame/methods/test_sort_values.py

@@ -657,32 +657,24 @@ def test_changes_length_raises(self):
         with pytest.raises(ValueError, match="change the shape"):
             df.sort_values("A", key=lambda x: x[:1])
 
-    def test_sort_values_key_dict(self):
+    def test_sort_values_key_axes(self):
         df = DataFrame({0: ["Hello", "goodbye"], 1: [0, 1]})
 
         result = df.sort_values(0, key=lambda col: col.str.lower())
         expected = df[::-1]
         tm.assert_frame_equal(result, expected)
 
-        result = df.sort_values([0, 1], key={0: lambda col: col.str.lower()})
+        result = df.sort_values(1, key=lambda col: -col)
         expected = df[::-1]
         tm.assert_frame_equal(result, expected)
 
-        result = df.sort_values([0, 1], key={1: lambda col: -col})
-        expected = df
-        tm.assert_frame_equal(result, expected)
-
     def test_sort_values_key_dict_axis(self):
         df = DataFrame({0: ["Hello", 0], 1: ["goodbye", 1]})
 
         result = df.sort_values(0, key=lambda col: col.str.lower(), axis=1)
         expected = df.loc[:, ::-1]
         tm.assert_frame_equal(result, expected)
 
-        result = df.sort_values([0, 1], key={0: lambda col: col.str.lower()}, axis=1)
+        result = df.sort_values(1, key=lambda col: -col, axis=1)
         expected = df.loc[:, ::-1]
         tm.assert_frame_equal(result, expected)
-
-        result = df.sort_values([0, 1], key={1: lambda col: -col}, axis=1)
-        expected = df
-        tm.assert_frame_equal(result, expected)