MultiRange: an implementation of multi ranges concatenated together in a Sequence interface

multi_range.py

from __future__ import annotations

import itertools
from abc import ABC
from collections import Sequence, deque

import typing
from overload import overload


class MultiRange(Sequence, ABC):
    ranges: typing.Deque[range]

    def __init__(self, ranges: Sequence[range] = None, pre_sorted: bool = False, seq_len: int = 1):
        """
        Create MultiRange with a sequence of ranges
        :param ranges: sequence of ranges
        :param pre_sorted: if the sequence is presorted
        :param seq_len: sequence_length
        """
        if ranges is None:
            ranges = []
        if not pre_sorted:
            ranges = sorted(ranges, key=lambda the_range: the_range.start)
        for i in range(len(ranges) - 1):
            assert ranges[i].stop <= ranges[i + 1].stop  # Multi range only support ranges without overlap
        self.ranges = deque(ranges)
        self.seq_len = seq_len

    @overload
    def __getitem__(self, i: int) -> int:
        range_id, item_id = self._get_item_id_index(i)
        return self.ranges[range_id][item_id]

    @__getitem__.add
    def __getitem__(self, s: slice) -> typing.Sequence[int]:
        new_ranges = []
        start_range_id, start_item_id = self._get_item_id_index(s.start)
        stop_range_id, stop_item_id = self._get_item_id_index(s.stop)
        for range_id, the_range in itertools.islice(enumerate(self.ranges), start_range_id, stop_range_id + 1):
            start = the_range.start if range_id != start_range_id else start_item_id
            stop = the_range.stop if range_id != stop_range_id else stop_item_id
            new_range = range(
                start,
                stop,
                s.step * the_range.step
            )
            new_ranges += [new_range]
        return MultiRange(new_ranges, pre_sorted=True)

    def __len__(self):
        return sum([len(temp_range) for temp_range in self.ranges])

    def _get_item_id_index(self, i: int) -> typing.Tuple[int, int]:
        """
        From an index of an item in the MultiRange to the index of the range that include that item and the index
        of the item inside the range
        :param i: Index of an item in the MultiRange
        :return: the index of the range that include that item and the index of the item inside the range
        """
        for range_id, the_range in enumerate(self.ranges):
            if i < len(the_range):
                return range_id, i
            else:
                i -= len(the_range)
        raise IndexError

    def _get_item_index(self, item: int) -> typing.Tuple[int, int]:
        """
        From an value of an item in the MultiRange to the index of the range that include that item and the index
        of the item inside the range
        :param item: Value of an item in the MultiRange
        :return: the index of the range that include that item and the index of the item inside the range
        """
        for range_id, the_range in enumerate(self.ranges):
            if item < the_range.start:
                break
            if item in range(the_range.start, the_range.stop + (self.seq_len - 1), the_range.step):
                return range_id, (item - the_range.start) // the_range.step
        raise IndexError

    def get_seq_index(self, seq_len: int) -> MultiRange:
        """
        Form a MultiRange that is the start index of a list of sequence with specified length that is parted of the
        current MultiRange
        :param seq_len: Length of each sequence
        :return: Generated MultiRange
        """
        new_ranges = [
            range(old_range.start, old_range.stop - (seq_len - 1), old_range.step)
            for old_range in self.ranges
            if old_range.stop - old_range.start >= seq_len
        ]
        return MultiRange(new_ranges, seq_len=seq_len)

    def _remove_forward(self, value, seq_len=None) -> int:  # Remove value but don't restore order in self.range
        """
        Remove the item with specified value without restore the order of the ranges
        :param value: Value of the item to be removed
        :param seq_len: Length of the sequence
        :return:
        """
        if seq_len is None:
            seq_len = self.seq_len
        range_id, _ = self._get_item_index(value)
        self.ranges.rotate(-range_id)
        old_range = self.ranges.popleft()
        new_ranges = [
            range(value + seq_len, old_range.stop, old_range.step),
            range(old_range.start, value - (seq_len - 1), old_range.step)
        ]
        new_existed_ranges = [the_range for the_range in new_ranges if len(the_range) > 0]  # remove empty ranges
        self.ranges.extendleft(new_existed_ranges)
        return range_id

    def remove(self, value, seq_len=None):
        """
        Remove the item with specified value
        :param value: Value of the item to be removed
        :param seq_len: Length of the sequence
        :return:
        """
        if seq_len is None:
            seq_len = self.seq_len
        range_id = self._remove_forward(value, seq_len=seq_len)
        self.ranges.rotate(range_id)

    def remove_items(self, items: typing.Sequence[int], seq_len: typing.Optional[int] = None, pre_sorted: bool = False):
        """
        Remove the items with specified value
        :param items: Values of the items to be removed
        :param seq_len: Length of the sequence
        :param pre_sorted: If the item values are pre-sorted
        :return:
        """
        if seq_len is None:
            seq_len = self.seq_len
        if not pre_sorted:
            items = sorted(items)
        total_rotation = 0
        for item in items:
            try:
                total_rotation += self._remove_forward(item, seq_len=seq_len)
            except IndexError:
                pass  # normal to have index error
        self.ranges.rotate(total_rotation)

    def __str__(self) -> str:
        return f"MultiRange[{', '.join([str(the_range) for the_range in self.ranges])}]"


if __name__ == '__main__':
    import copy

    mr = MultiRange([range(1, 5), range(10, 13), range(7, 9), range(20, 31)])
    print(f"Original MultiRange: {mr}")
    print(f"Original MultiRange items: {list(mr)}")
    mr_removed = copy.deepcopy(mr)
    mr_removed.remove(25)
    print(f"Original MultiRange with 25 removed: {mr_removed}")
    mr_removed.remove_items([2, 8, 21, 27])
    print(f"Last MultiRange with 2, 8, 21, 27 removed: {mr_removed}")
    mr_2 = mr.get_seq_index(2)
    print(f"Start index as sequence of 2: {mr_2}")
    mr_2_removed = copy.deepcopy(mr_2)
    mr_2_removed.remove(25)
    print(f"Index MultiRange with 25 removed: {mr_2_removed}")
    mr_2_removed.remove_items([2, 8, 21, 27])
    print(f"Last MultiRange with 2, 8, 21, 27 removed: {mr_2_removed}")