Source code

Revision control

Copy as Markdown

Other Tools

""" 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