def _ensure_key_mapped_multiindex(
index: MultiIndex, key: Callable, level=None
) -> MultiIndex:
"""
Returns a new MultiIndex in which key has been applied
to all levels specified in level (or all levels if level
is None). Used for key sorting for MultiIndex.
Parameters
----------
index : MultiIndex
Index to which to apply the key function on the
specified levels.
key : Callable
Function that takes an Index and returns an Index of
the same shape. This key is applied to each level
separately. The name of the level can be used to
distinguish different levels for application.
level : 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
levels are left unchanged.
Returns
-------
labels : MultiIndex
Resulting MultiIndex with modified levels.
"""
if level is not None:
if isinstance(level, (str, int)):
sort_levels = [level]
else:
sort_levels = level
sort_levels = [index._get_level_number(lev) for lev in sort_levels]
else:
sort_levels = list(range(index.nlevels)) # satisfies mypy
mapped = [
ensure_key_mapped(index._get_level_values(level), key)
if level in sort_levels
else index._get_level_values(level)
for level in range(index.nlevels)
]
return type(index).from_arrays(mapped)
