Ejemplo n.º 1
0
 def test_loc_map_b(self) -> None:
     idx = Index(['a', 'b', 'c', 'd', 'e'])
     post1 = LocMap.loc_to_iloc(
         label_to_pos=idx._map,
         labels=idx._labels,
         positions=idx._positions,
         key=['b', 'd'],
         partial_selection=False,
     )
     self.assertEqual(post1, [1, 3])
Ejemplo n.º 2
0
    def test_loc_map_a(self) -> None:
        idx = Index(['a', 'b', 'c'])
        post1 = LocMap.loc_to_iloc(
            label_to_pos=idx._map,
            labels=idx._labels,
            positions=idx._positions,
            key='b',
            partial_selection=False,
        )
        self.assertEqual(post1, 1)

        post2 = LocMap.loc_to_iloc(
            label_to_pos=idx._map,
            labels=idx._labels,
            positions=idx._positions,
            key=NULL_SLICE,
            partial_selection=False,
        )
        self.assertEqual(post2, NULL_SLICE)
Ejemplo n.º 3
0
    def test_loc_map_slice_c(self) -> None:
        dt64 = np.datetime64
        idx = IndexDate.from_date_range('1985-01-01', '1985-01-08')

        post1 = LocMap.loc_to_iloc(
            label_to_pos=idx._map,
            labels=idx._labels,
            positions=idx._positions,
            key=slice(dt64('1985-01-01'), dt64('1985-01-04')),
            partial_selection=False,
        )
        self.assertEqual(post1, slice(0, 4, None))
Ejemplo n.º 4
0
    def test_loc_map_slice_b(self) -> None:
        dt64 = np.datetime64
        idx = IndexDate.from_date_range('1985-01-01', '1985-01-08')

        with self.assertRaises(RuntimeError):
            post1 = LocMap.loc_to_iloc(
                label_to_pos=idx._map,
                labels=idx._labels,
                positions=idx._positions,
                key=slice(dt64('1985-01-01'), dt64('1985-01-04'),
                          dt64('1985-01-04')),
                partial_selection=False,
            )
Ejemplo n.º 5
0
    def loc_to_iloc(self, key: GetItemKeyTypeCompound) -> GetItemKeyType:
        '''
        This is the low-level loc_to_iloc, analagous to LocMap.loc_to_iloc as used by Index. As such, the key at this point should not be a Series or Index object.

        If key is an np.ndarray, a Boolean array will be passed through; otherwise, it will be treated as an iterable of values to be passed to leaf_loc_to_iloc.
        '''
        from static_frame.core.series import Series

        if isinstance(key, slice):
            # given a top-level definition of a slice (and if that slice results in a single value), we can get a value range
            return slice(*LocMap.map_slice_args(self.leaf_loc_to_iloc, key))

        if isinstance(key, KEY_ITERABLE_TYPES):  # iterables of leaf-locs
            if key.__class__ is np.ndarray and key.dtype == bool:  #type: ignore
                return key  # keep as Boolean
            return [self.leaf_loc_to_iloc(x) for x in key]

        if not isinstance(key, HLoc):  # assume a leaf loc tuple
            if not isinstance(key, tuple):
                raise KeyError(
                    f'{key} cannot be used for loc selection from IndexHierarchy; try HLoc'
                )
            return self.leaf_loc_to_iloc(key)

        # HLoc following: collect all ilocs for all leaf indices matching HLoc patterns
        ilocs = []
        levels = deque(((self, 0, 0), ))  # order matters

        while levels:
            level, depth, offset = levels.popleft()

            depth_key = key[depth]
            # NOTE: depth_key should not be Series or Index at this point; IndexHierarchy is responsible for unpacking / reindexing prior to this call
            next_offset = offset + level.offset

            if depth_key.__class__ is np.ndarray and depth_key.dtype == DTYPE_BOOL:  #type: ignore
                # NOTE: use length of level, not length of index, as need to observe all leafs covered at this node.
                depth_key = depth_key[next_offset:next_offset +
                                      len(level)]  #type: ignore
                if len(depth_key) > len(level.index):  #type: ignore
                    # given leaf-Boolean, determine what upper nodes to select
                    depth_key = level.values_at_depth(0)[depth_key]
                    if len(depth_key) > 1:
                        # NOTE: must strip repeated labels, but cannot us np.unique as must retain order
                        depth_key = list(dict.fromkeys(depth_key).keys())

            # print(level, depth, offset, depth_key, next_offset)
            if level.targets is None:
                try:
                    # NOTE: as a selection list might be given within the HLoc, it will be tested accross many indices, and should support a partial matching
                    ilocs.append(
                        level.index._loc_to_iloc(
                            depth_key,
                            offset=next_offset,
                            partial_selection=True,
                        ))
                except KeyError:
                    pass
            else:  # when not at a leaf, we are selecting level_targets to descend withing
                try:  # NOTE: no offset necessary as not a leaf selection
                    iloc = level.index._loc_to_iloc(depth_key,
                                                    partial_selection=True)
                except KeyError:
                    pass
                else:
                    level_targets = level.targets[
                        iloc]  # get one or more IndexLevel objects
                    next_depth = depth + 1
                    # if not an ndarray, iloc has extracted a single IndexLevel
                    if isinstance(level_targets, IndexLevel):
                        levels.append((level_targets, next_depth, next_offset))
                    else:
                        levels.extend((lvl, next_depth, next_offset)
                                      for lvl in level_targets)

        iloc_count = len(ilocs)
        if iloc_count == 0:
            raise KeyError('no matching keys across all levels')

        if iloc_count == 1 and not key.has_key_multiple():
            return ilocs[0]  # drop to a single iloc selection

        # NOTE: might be able to combine contiguous ilocs into a single slice
        iloc_flat: tp.List[GetItemKeyType] = []  # combine into one flat iloc
        length = self.__len__()
        for part in ilocs:
            if isinstance(part, slice):
                iloc_flat.extend(range(*part.indices(length)))
            elif isinstance(part, INT_TYPES):
                iloc_flat.append(part)
            else:  # assume it is an iterable
                iloc_flat.extend(part)  #type: ignore
        return iloc_flat