SDSM-for-SDI / src / aioquic / quic / rangeset.py
rangeset.py
Raw
from collections.abc import Sequence
from typing import Any, Iterable, List, Optional


class RangeSet(Sequence):
    def __init__(self, ranges: Iterable[range] = []):
        self.__ranges: List[range] = []
        for r in ranges:
            assert r.step == 1
            self.add(r.start, r.stop)

    def add(self, start: int, stop: Optional[int] = None) -> None:
        if stop is None:
            stop = start + 1
        assert stop > start

        for i, r in enumerate(self.__ranges):
            # the added range is entirely before current item, insert here
            if stop < r.start:
                self.__ranges.insert(i, range(start, stop))
                return

            # the added range is entirely after current item, keep looking
            if start > r.stop:
                continue

            # the added range touches the current item, merge it
            start = min(start, r.start)
            stop = max(stop, r.stop)
            while i < len(self.__ranges) - 1 and self.__ranges[i + 1].start <= stop:
                stop = max(self.__ranges[i + 1].stop, stop)
                self.__ranges.pop(i + 1)
            self.__ranges[i] = range(start, stop)
            return

        # the added range is entirely after all existing items, append it
        self.__ranges.append(range(start, stop))

    def bounds(self) -> range:
        return range(self.__ranges[0].start, self.__ranges[-1].stop)

    def shift(self) -> range:
        return self.__ranges.pop(0)

    def subtract(self, start: int, stop: int) -> None:
        assert stop > start

        i = 0
        while i < len(self.__ranges):
            r = self.__ranges[i]

            # the removed range is entirely before current item, stop here
            if stop <= r.start:
                return

            # the removed range is entirely after current item, keep looking
            if start >= r.stop:
                i += 1
                continue

            # the removed range completely covers the current item, remove it
            if start <= r.start and stop >= r.stop:
                self.__ranges.pop(i)
                continue

            # the removed range touches the current item
            if start > r.start:
                self.__ranges[i] = range(r.start, start)
                if stop < r.stop:
                    self.__ranges.insert(i + 1, range(stop, r.stop))
            else:
                self.__ranges[i] = range(stop, r.stop)

            i += 1

    def __bool__(self) -> bool:
        raise NotImplementedError

    def __contains__(self, val: Any) -> bool:
        for r in self.__ranges:
            if val in r:
                return True
        return False

    def __eq__(self, other: object) -> bool:
        if not isinstance(other, RangeSet):
            return NotImplemented

        return self.__ranges == other.__ranges

    def __getitem__(self, key: Any) -> range:
        return self.__ranges[key]

    def __len__(self) -> int:
        return len(self.__ranges)

    def __repr__(self) -> str:
        return "RangeSet({})".format(repr(self.__ranges))