def preprocess_index()

in core/maxframe/tensor/indexing/core.py [0:0]

• 53 lines of code
• 28 McCabe index (conditional complexity)

def preprocess_index(index, convert_bool_to_fancy=None):
    from .nonzero import nonzero

    inds = []
    fancy_indexes = []
    bool_indexes = []
    all_fancy_index_ndarray = True
    all_bool_index_ndarray = True
    for j, ind in enumerate(index):
        if isinstance(ind, (list, np.ndarray) + TENSOR_TYPE):
            if not isinstance(ind, TENSOR_TYPE):
                ind = np.array(ind)
            if ind.dtype.kind not in "biu":
                raise IndexError(_INDEX_ERROR_MSG)
            if ind.dtype.kind == "b":
                # bool indexing
                bool_indexes.append(j)
                if not isinstance(ind, np.ndarray):
                    all_bool_index_ndarray = False
            else:
                # fancy indexing
                fancy_indexes.append(j)
                if not isinstance(ind, np.ndarray):
                    all_fancy_index_ndarray = False
        elif (
            not isinstance(ind, (slice, Integral))
            and ind is not None
            and ind is not Ellipsis
        ):
            raise IndexError(_INDEX_ERROR_MSG)
        inds.append(ind)

    if convert_bool_to_fancy is None:
        convert_bool_to_fancy = (fancy_indexes and len(bool_indexes) > 0) or len(
            bool_indexes
        ) > 1

    if not all_fancy_index_ndarray or (
        convert_bool_to_fancy and not all_bool_index_ndarray
    ):
        # if not all fancy indexes are ndarray,
        # or bool indexes need to be converted to fancy indexes,
        # and not all bool indexes are ndarray,
        # we will convert all of them to Tensor
        for fancy_index in fancy_indexes:
            inds[fancy_index] = astensor(inds[fancy_index])

    # convert bool index to fancy index when any situation below meets:
    # 1. fancy indexes and bool indexes both exists
    # 2. bool indexes more than 2
    if convert_bool_to_fancy:
        default_m = None
        if len(fancy_indexes) > 0:
            default_m = (
                np.nonzero
                if isinstance(inds[fancy_indexes[0]], np.ndarray)
                else nonzero
            )
        for bool_index in bool_indexes:
            ind = inds[bool_index]
            m = default_m
            if m is None:
                m = np.nonzero if isinstance(ind, np.ndarray) else nonzero
            ind = m(ind)[0]
            inds[bool_index] = ind

    return tuple(inds)