Quelle  ordered.py   Sprache: Python

""" Deterministic data structures. """

from __future__ import annotations
# mypy: disallow-untyped-defs, disallow-incomplete-defs, disallow-untyped-calls

from typing import (AbstractSet, Dict, Generic, Iterable, Iterator, List,
                    MutableSet, Optional, TypeVar, Union)


__all__ = ['OrderedSet', 'OrderedFrozenSet']


# Type parameters for OrderedSet[T] and OrderedFrozenSet[T].
#
# These should be `bound=Hashable`, but I gather MyPy has some issues with
# hashability. It doesn't enforce hashability on Dict and Set.
T = TypeVar('T')
T_co = TypeVar('T_co', covariant=True)

S = TypeVar('S')


class OrderedSet(Generic[T], MutableSet[T]):
    """Like set(), but iteration order is insertion order.

    Two OrderedSets, x and y, that have different insertion order are still
    considered equal (x == y) if they contain the same elements.
    """
    _data: Dict[T, int]

    def __init__(self, values: Iterable[T] = ()):
        self._data = {}
        for v in values:
            self.add(v)

    def __repr__(self) -> str:
        return self.__class__.__name__ + "(" + repr(list(self)) + ")"

    def add(self, v: T) -> None:
        self._data[v] = 1

    def extend(self, iterable: Iterable[T]) -> None:
        for v in iterable:
            self.add(v)

    def remove(self, v: T) -> None:
        del self._data[v]

    def discard(self, v: T) -> None:
        if v in self._data:
            del self._data[v]

    def __eq__(self, other: object) -> bool:
        return isinstance(other, OrderedSet) and self._data == other._data

    def __hash__(self) -> int:
        raise TypeError("unhashable type: " + self.__class__.__name__)

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

    def __bool__(self) -> bool:
        return bool(self._data)

    def __contains__(self, v: object) -> bool:
        return v in self._data

    def __iter__(self) -> Iterator[T]:
        return iter(self._data)

    def __ior__(self, other: AbstractSet[S]) -> OrderedSet[Union[T, S]]:
        for v in other:
            self.add(v)  # type: ignore
        return self  # type: ignore

    def __or__(self, other: AbstractSet[S]) -> OrderedSet[Union[T, S]]:
        u: OrderedSet[Union[T, S]] = OrderedSet(self)
        u |= other
        return u

    def __iand__(self, other: AbstractSet[T]) -> OrderedSet[T]:
        self._data = {v: 1 for v in self if v in other}
        return self

    def __and__(self, other: AbstractSet[T]) -> OrderedSet[T]:
        return OrderedSet(v for v in self if v in other)

    def __sub__(self, other: AbstractSet[T]) -> OrderedSet[T]:
        return OrderedSet(v for v in self if v not in other)

    def __isub__(self, other: AbstractSet[T]) -> OrderedSet[T]:
        for v in other:
            if v in self:
                self.remove(v)
        return self

    def is_disjoint(self, other: AbstractSet[T]) -> bool:
        for v in self:
            if v in other:
                return False
        return True


class OrderedFrozenSet(Generic[T_co], AbstractSet[T_co]):
    """Like frozenset(), but iteration order is insertion order.

    Two OrderedFrozenSets, x and y, that have different insertion order are
    still considered equal (x == y) if they contain the same elements.
    """
    __slots__ = ['_data', '_hash']

    _data: Dict[T_co, int]
    _hash: Optional[int]

    def __init__(self, values: Iterable[T_co] = ()):
        self._data = {v: 1 for v in values}
        self._hash = None

    def __repr__(self) -> str:
        return self.__class__.__name__ + "(" + repr(list(self)) + ")"

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

    def __bool__(self) -> bool:
        return bool(self._data)

    def __contains__(self, v: object) -> bool:
        return v in self._data

    def __iter__(self) -> Iterator[T_co]:
        return iter(self._data)

    def __eq__(self, other: object) -> bool:
        return isinstance(other, OrderedFrozenSet) and self._data == other._data

    def __hash__(self) -> int:
        if self._hash is None:
            self._hash = hash(frozenset(self._data))
        return self._hash

    def __and__(self, other: AbstractSet[T_co]) -> OrderedFrozenSet[T_co]:
        return OrderedFrozenSet(v for v in self._data if v in other)

    def __or__(self, other: AbstractSet[S]) -> OrderedFrozenSet[Union[T_co, S]]:
        values: List[Union[T_co, S]] = list(self)
        values += list(other)
        return OrderedFrozenSet(values)

    def __sub__(self, other: AbstractSet[T_co]) -> OrderedFrozenSet[T_co]:
        return OrderedFrozenSet(v for v in self._data if v not in other)

    def is_disjoint(self, other: AbstractSet[T_co]) -> bool:
        for v in self:
            if v in other:
                return False
        return True