Exemplo n.º 1
0
def add_index_level(
    old_index: pandas.Index, value: Any, name: str = None, loc: int = 0
) -> pandas.MultiIndex:
    """
    Expand a (multi)index by adding a level to it.

    :param old_index: The index to expand
    :param name: The name of the new index level
    :param value: Scalar or list-like, the values of the new index level
    :param loc: Where to insert the level in the index, 0 is at the front, negative values count back from the rear end
    :return: A new multi-index with the new level added
    """

    def _handle_insert_loc(loc: int, n: int) -> int:
        """
        Computes the insert index from the right if loc is negative for a given size of n.
        """
        return n + loc + 1 if loc < 0 else loc

    loc = _handle_insert_loc(loc, len(old_index.names))
    old_index_df = old_index.to_frame()
    old_index_df.insert(loc, name, value)
    new_index_names = list(
        old_index.names
    )  # sometimes new index level names are invented when converting to a df,
    new_index_names.insert(loc, name)  # here the original names are reconstructed
    new_index = pandas.MultiIndex.from_frame(old_index_df, names=new_index_names)
    return new_index
Exemplo n.º 2
0
    def _get_indexer(self, index: pd.Index) -> Iterable[bool]:
        if not callable(self.cond):
            return self.cond

        func = partial(self.cond, **self.kwargs)

        if isinstance(index, pd.MultiIndex):
            mi_df = index.to_frame()
            selected = mi_df[Everywhere(func)]
            mi_values = pd.MultiIndex.from_frame(selected).to_numpy()
            return index.get_locs(mi_values)

        return func(index)