import itertools
import itertools as it
from typing import (
AbstractSet,
Any,
Dict,
Iterable,
Iterator,
List,
MutableSet,
Optional,
Sequence,
Set,
TypeVar,
Union,
overload,
)
SLICE_ALL = slice(None)
T = TypeVar("T")
# SetLike[T] is either a set of elements of type T, or a sequence, which
# we will convert to a StableSet or to an OrderedSet by adding its elements in order.
SetLike = Union[AbstractSet[T], Sequence[T]]
SetInitializer = Union[AbstractSet[T], Sequence[T], Iterable[T]]
def _is_atomic(obj: object) -> bool:
"""
Returns True for objects which are iterable but should not be iterated in
the context of indexing a StableSet or an OrderedSet.
When we index by an iterable, usually that means we're being asked to look
up a list of things.
However, in the case of the .index() method, we shouldn't handle strings
and tuples like other iterables. They're not sequences of things to look
up, they're the single, atomic thing we're trying to find.
As an example, oset.index('hello') should give the index of 'hello' in an
StableSet of strings. It shouldn't give the indexes of each individual
character.
"""
return isinstance(obj, (str, tuple))
class StableSet(MutableSet[T], Sequence[T]):
"""
A StableSet is a custom MutableSet that remembers its insertion order.
Featuring: Fast O(1) insertion, deletion, iteration and membership testing.
But slow O(N) Index Lookup.
StableSet is meant to be a drop-in replacement for `set` when iteration in insertion order
is the only additional requirement over the built-in `set`.
Equality: StableSet, like `set` and `dict_keys` [dict.keys()], and unlike OrderdSet,
disregards the items order when checking equality.
Like `set` it may be equal only to other instances of AbstractSet
(like `set`, `dict_keys` or StableSet).
This implementation of StableSet is based on the built-in dict type.
In Python 3.6 and later, the built-in dict type is inherently ordered.
If you ignore the dictionary values, that also gives you a simple ordered set,
with fast O(1) insertion, deletion, iteration and membership testing.
However, dict does not provide the list-like random access features of StableSet.
So we have to convert it to a list in O(N) to look up the index of an entry
or look up an entry by its index.
Example:
>>> StableSet([1, 1, 2, 3, 2])
StableSet([1, 2, 3])
"""
__slots__ = ("_map", "_is_mutable")
_map: Dict[T, Any]
def __init__(self, initial: Optional[SetInitializer[T]] = None):
self._map = dict.fromkeys(initial) if initial else {}
self._is_mutable = True
def __len__(self) -> int:
"""
Returns the number of unique elements in the ordered set
Example:
>>> len(StableSet([]))
0
>>> len(StableSet([1, 2]))
2
"""
return self._map.__len__()
@overload
def __getitem__(self, index: slice) -> "StableSet[T]":
...
@overload
def __getitem__(self, index: Sequence[int]) -> List[T]:
...
@overload
def __getitem__(self, index: int) -> T:
...
# concrete implementation
def __getitem__(self, index):
"""
Get the item at a given index.
If `index` is a slice, you will get back that slice of items, as a
new StableSet.
If `index` is a list or a similar iterable, you'll get a list of
items corresponding to those indices. This is similar to NumPy's
"fancy indexing". The result is not a StableSet because you may ask
for duplicate indices, and the number of elements returned should be
the number of elements asked for.
Example:
>>> oset = StableSet([1, 2, 3])
>>> oset[1]
2
"""
if isinstance(index, int):
if index < 0:
index = len(self._map) + index
try:
return next(itertools.islice(self._map.keys(), index, index + 1))
except StopIteration:
raise IndexError(f"index {index} out of range")
elif isinstance(index, slice) and index == SLICE_ALL:
return self.copy()
items = list(self._map.keys())
if isinstance(index, Iterable):
return [items[i] for i in index]
elif isinstance(index, slice) or hasattr(index, "__index__"):
result = items[index]
if isinstance(result, list):
return self.__class__(result)
else:
return result
else:
raise TypeError(f"Don't know how to index a StableSet by {index}")
# Define the gritty details of how a StableSet is serialized as a pickle.
# We leave off type annotations, because the only code that should interact
# with this is a generalized tool such as pickle.
def __getstate__(self):
if len(self) == 0:
# In pickle, the state can't be an empty list.
# We need to return a truthy value, or else __setstate__ won't be run.
#
# This could have been done more gracefully by always putting the state
# in a tuple, but this way is backwards- and forwards- compatible with
# previous versions of StableSet.
return (None,)
else:
return list(self)
def __setstate__(self, state):
if state == (None,):
self.__init__([])
else:
self.__init__(state)
def __contains__(self, key: Any) -> bool:
"""
Test if the item is in this ordered set.
Example:
>>> 1 in StableSet([1, 3, 2])
True
>>> 5 in StableSet([1, 3, 2])
False
"""
# return key in self._map
return self._map.__contains__(key)
def __iter__(self) -> Iterator[T]:
"""
Example:
>>> list(iter(StableSet([1, 2, 3])))
[1, 2, 3]
"""
# return iter(self._map.keys())
return self._map.keys().__iter__()
def __reversed__(self) -> Iterator[T]:
"""
Supported from Python >= 3.8
Example:
>>> list(reversed(StableSet([1, 2, 3])))
[3, 2, 1]
"""
return reversed(self._map.keys())
def __repr__(self) -> str:
if not self:
return f"{self.__class__.__name__}()"
return f"{self.__class__.__name__}({list(self)!r})"
__str__ = __repr__
def __and__(self, other: SetLike[T]) -> "StableSet[T]":
# the parent implementation of this is backwards
return self.intersection(other)
# sub, or, xor that support ordering
# (left hand and right hand - as the operands order does matter)
# based on the implementations of the super class (Set(Collection)),
# see _collections_abc.py
def __sub__(self, other):
cls = type(
self
if isinstance(self, StableSet)
else other
if isinstance(other, StableSet)
else StableSet
)
if not isinstance(other, Set):
if not isinstance(other, Iterable):
return NotImplemented
other = cls(other)
return cls(value for value in self if value not in other)
def __rsub__(self, other):
cls = type(
self
if isinstance(self, StableSet)
else other
if isinstance(other, StableSet)
else StableSet
)
if not isinstance(other, Set):
if not isinstance(other, Iterable):
return NotImplemented
other = cls(other)
return cls(value for value in other if value not in self)
def __or__(self, other):
cls = type(
self
if isinstance(self, StableSet)
else other
if isinstance(other, StableSet)
else StableSet
)
if not isinstance(other, Iterable):
return NotImplemented
chain = (e for s in (self, other) for e in s)
return cls(chain)
def __ror__(self, other):
cls = type(
self
if isinstance(self, StableSet)
else other
if isinstance(other, StableSet)
else StableSet
)
if not isinstance(other, Iterable):
return NotImplemented
chain = (e for s in (other, self) for e in s)
return cls(chain)
def __xor__(self, other):
if not isinstance(other, Iterable):
return NotImplemented
return (self - other) | (other - self)
def __rxor__(self, other):
if not isinstance(other, Iterable):
return NotImplemented
return (other - self) | (self - other)
def __eq__(self, other):
if not isinstance(other, Iterable):
return False
if len(self._map) != len(other):
return False
if isinstance(other, StableSet):
return self._map == other._map
if not isinstance(other, list):
other = list(other)
return list(self._map.keys()) == other
def clear(self) -> None:
"""
Remove all items from this StableSet.
"""
if self._is_mutable is False:
raise ValueError("This object is not mutable.")
self._map.clear()
def copy(self) -> "StableSet[T]":
"""
Return a shallow copy of this object.
Example:
>>> this = StableSet([1, 2, 3])
>>> other = this.copy()
>>> this == other
True
>>> this is other
False
"""
return self.__class__(self)
# Technically type-incompatible with MutableSet, because we return an
# int instead of nothing. This is also one of the things that makes
# StableSet convenient to use.
def add(self, key: T) -> int:
"""
Add `key` as an item to this StableSet, then return its index.
If `key` is already in the StableSet, return the index it already
had.
Example:
>>> oset = StableSet()
>>> oset.append(3)
0
>>> print(oset)
StableSet([3])
"""
if self._is_mutable is False:
raise ValueError("This object is not mutable.")
self._map[key] = None
return len(self._map) - 1
append = add
def update(self, sequence: SetLike[T]) -> int:
"""
Update the set with the given iterable sequence, then return the index
of the last element inserted.
Example:
>>> oset = StableSet([1, 2, 3])
>>> oset.update([3, 1, 5, 1, 4])
4
>>> print(oset)
StableSet([1, 2, 3, 5, 4])
"""
if self._is_mutable is False:
raise ValueError("This object is not mutable.")
other_map = dict.fromkeys(sequence)
self._map.update(other_map)
return len(self._map) - 1
@overload
def index(self, key: Sequence[T]) -> List[int]: # NOQA
...
@overload
def index(self, key: T) -> int: # NOQA
...
# concrete implementation
def index(self, key): # NOQA
"""
Get the index of a given entry, raising an IndexError if it's not present
`key` can be an iterable of entries that is not a string, in which case
this returns a list of indices.
Example:
>>> oset = StableSet([1, 2, 3])
>>> oset.index(2)
1
"""
try:
if isinstance(key, Iterable) and not _is_atomic(key):
return [self.index(subkey) for subkey in key]
for index, item in enumerate(self._map.keys()):
if item == key:
return index
raise KeyError(key)
# return list(self._map.keys()).index(key)
except ValueError:
raise KeyError(key)
# Provide some compatibility with pd.Index
get_loc = index
get_indexer = index
def pop(self, index: int = -1) -> T:
"""
Remove and return item at index (default last).
Raises KeyError if the set is empty.
Raises IndexError if index is out of range.
Example:
>>> oset = StableSet([1, 2, 3])
>>> oset.pop()
3
"""
if self._is_mutable is False:
raise ValueError("This object is not mutable.")
if not self._map:
raise KeyError("Set is empty")
if index == -1:
elem, _ = self._map.popitem()
return elem
elif index == 0:
elem = next(iter(self._map.keys()))
else:
elem = next(itertools.islice(self._map.keys(), index, index + 1))
self._map.pop(elem)
return elem
def popitem(self, last: bool = True):
"""Remove and return an item from the set.
Items are returned in LIFO order if last is true or FIFO order if false.
"""
if self._is_mutable is False:
raise ValueError("This object is not mutable.")
if not self._map:
raise KeyError("Set is empty")
if last:
elem, _ = self._map.popitem()
return elem
elem = next(iter(self._map.keys()))
self._map.pop(elem)
return elem
def move_to_end(self, key) -> None:
"""Move an existing element to the end.
Raise KeyError if the element does not exist.
"""
if self._is_mutable is False:
raise ValueError("This object is not mutable.")
self._map.pop(key)
self._map[key] = None
def discard(self, key: T) -> None:
"""
Remove an element. Do not raise an exception if absent.
The MutableSet mixin uses this to implement the .remove() method, which
*does* raise an error when asked to remove a non-existent item.
Example:
>>> oset = StableSet([1, 2, 3])
>>> oset.discard(2)
>>> print(oset)
StableSet([1, 3])
>>> oset.discard(2)
>>> print(oset)
StableSet([1, 3])
"""
if self._is_mutable is False:
raise ValueError("This object is not mutable.")
self._map.pop(key, None)
def union(self, *sets: SetLike[T]) -> "StableSet[T]":
"""
Combines all unique items.
Each item order is defined by its first appearance.
Example:
>>> oset = StableSet.union(StableSet([3, 1, 4, 1, 5]), [1, 3], [2, 0])
>>> print(oset)
StableSet([3, 1, 4, 5, 2, 0])
>>> oset.union([8, 9])
StableSet([3, 1, 4, 5, 2, 0, 8, 9])
>>> oset | {10}
StableSet([3, 1, 4, 5, 2, 0, 10])
"""
cls = type(self if isinstance(self, StableSet) else StableSet)
containers = map(list, it.chain([self], sets)) # type: ignore
items = it.chain.from_iterable(containers)
return cls(items) # type: ignore
def intersection(self, *sets: SetLike[T]) -> "StableSet[T]":
"""
Returns elements in common between all sets. Order is defined only
by the first set.
Example:
>>> oset = StableSet.intersection(StableSet([0, 1, 2, 3]), [1, 2, 3])
>>> print(oset)
StableSet([1, 2, 3])
>>> oset.intersection([2, 4, 5], [1, 2, 3, 4])
StableSet([2])
>>> oset.intersection()
StableSet([1, 2, 3])
"""
cls = type(self if isinstance(self, StableSet) else StableSet)
items: SetInitializer[T] = self
if sets:
common = set.intersection(*map(set, sets)) # type: ignore
items = (item for item in self if item in common)
return cls(items)
def difference(self, *sets: SetLike[T]) -> "StableSet[T]":
"""
Returns all elements that are in this set but not the others.
Example:
>>> StableSet([1, 2, 3]).difference(StableSet([2]))
StableSet([1, 3])
>>> StableSet([1, 2, 3]).difference(StableSet([2]), StableSet([3]))
StableSet([1])
>>> StableSet([1, 2, 3]) - StableSet([2])
StableSet([1, 3])
>>> StableSet([1, 2, 3]).difference()
StableSet([1, 2, 3])
"""
cls = type(self if isinstance(self, StableSet) else StableSet)
items: SetInitializer[T] = self
if sets:
other = set.union(*map(set, sets)) # type: ignore
items = (item for item in self if item not in other)
return cls(items)
def symmetric_difference(self, other: SetLike[T]) -> "StableSet[T]":
"""
Return the symmetric difference of two StableSets as a new set.
That is, the new set will contain all elements that are in exactly
one of the sets.
Their order will be preserved, with elements from `self` preceding
elements from `other`.
Example:
>>> this = StableSet([1, 4, 3, 5, 7])
>>> other = StableSet([9, 7, 1, 3, 2])
>>> this.symmetric_difference(other)
StableSet([4, 5, 9, 2])
"""
cls = type(
self
if isinstance(self, StableSet)
else other
if isinstance(other, StableSet)
else StableSet
)
diff1 = cls(self).difference(other)
diff2 = cls(other).difference(self)
return diff1.union(diff2)
def difference_update(self, *sets: SetLike[T]) -> None:
"""
Update this StableSet to remove items from one or more other sets.
Example:
>>> this = StableSet([1, 2, 3])
>>> this.difference_update(StableSet([2, 4]))
>>> print(this)
StableSet([1, 3])
>>> this = StableSet([1, 2, 3, 4, 5])
>>> this.difference_update(StableSet([2, 4]), StableSet([1, 4, 6]))
>>> print(this)
StableSet([3, 5])
"""
if self._is_mutable is False:
raise ValueError("This object is not mutable.")
items_to_remove = set() # type: Set[T]
for other in sets:
items_as_set = set(other) # type: Set[T]
items_to_remove |= items_as_set
self._map = dict.fromkeys(
[item for item in self._map if item not in items_to_remove]
)
def intersection_update(self, other: SetLike[T]) -> None:
"""
Update this StableSet to keep only items in another set, preserving
their order in this set.
Example:
>>> this = StableSet([1, 4, 3, 5, 7])
>>> other = StableSet([9, 7, 1, 3, 2])
>>> this.intersection_update(other)
>>> print(this)
StableSet([1, 3, 7])
"""
if self._is_mutable is False:
raise ValueError("This object is not mutable.")
other = set(other)
self._map = dict.fromkeys([item for item in self._map if item in other])
def symmetric_difference_update(self, other: SetLike[T]) -> None:
"""
Update this StableSet to remove items from another set, then
add items from the other set that were not present in this set.
Example:
>>> this = StableSet([1, 4, 3, 5, 7])
>>> other = StableSet([9, 7, 1, 3, 2])
>>> this.symmetric_difference_update(other)
>>> print(this)
StableSet([4, 5, 9, 2])
"""
items_to_add = [item for item in other if item not in self]
items_to_remove = set(other)
self._map = dict.fromkeys(
[item for item in self._map if item not in items_to_remove] + items_to_add
)
def issubset(self, other: SetLike[T]) -> bool:
"""
Report whether another set contains this set.
Example:
>>> StableSet([1, 2, 3]).issubset({1, 2})
False
>>> StableSet([1, 2, 3]).issubset({1, 2, 3, 4})
True
>>> StableSet([1, 2, 3]).issubset({1, 4, 3, 5})
False
"""
if len(self) > len(other): # Fast check for obvious cases
return False
return all(item in other for item in self)
def issuperset(self, other: SetLike[T]) -> bool:
"""
Report whether this set contains another set.
Example:
>>> StableSet([1, 2]).issuperset([1, 2, 3])
False
>>> StableSet([1, 2, 3, 4]).issuperset({1, 2, 3})
True
>>> StableSet([1, 4, 3, 5]).issuperset({1, 2, 3})
False
"""
if len(self) < len(other): # Fast check for obvious cases
return False
return all(item in self for item in other)
def isorderedsubset(self: SetLike, other: SetLike, non_consecutive: bool = False):
if len(self) > len(other):
return False
if non_consecutive:
i = 0
self_len = len(self)
for other_item in other:
if other_item == self[i]:
i += 1
if i == self_len:
return True
return False
else:
for self_item, other_item in zip(self, other):
if not self_item == other_item:
return False
return True
def isorderedsuperset(self, other: SetLike, non_consecutive: bool = False):
return StableSet.isorderedsubset(other, self, non_consecutive)
def get(self):
return next(iter(self._map))
def freeze(self):
"""
Once this function is run, the object becomes immutable
"""
self._is_mutable = False
class OrderlySet(StableSet[T]):
"""
OrderlySet keeps the order when adding but if you do difference, subtraction, etc, you lose the order.
The new results will have a random order but they will keep that order.
"""
def __sub__(self, other):
other = other if isinstance(other, (set, frozenset)) else set(other)
result = set(self) - other
return OrderlySet(result)
def __rsub__(self, other):
other = other if isinstance(other, (set, frozenset)) else set(other)
result = other - set(self)
return OrderlySet(result)
def __xor__(self, other):
other = other if isinstance(other, (set, frozenset)) else set(other)
result = set(self) ^ other
return OrderlySet(result)
__rxor__ = __xor__
def __eq__(self, other):
if not isinstance(other, Iterable):
return False
if len(self._map) != len(other):
return False
if isinstance(other, StableSet):
return self._map == other._map
if not isinstance(other, (set, frozenset)):
other = set(other)
return set(self._map.keys()) == other
def __ge__(self, other):
if not isinstance(other, Iterable):
return False
if len(self._map) < len(other):
return False
if not isinstance(other, (set, frozenset)):
other = set(other)
return set(self._map.keys()) >= other
def __gt__(self, other):
if not isinstance(other, Iterable):
return False
if len(self._map) <= len(other):
return False
if not isinstance(other, (set, frozenset)):
other = set(other)
return set(self._map.keys()) > other
def __le__(self, other):
if not isinstance(other, Iterable):
return False
if len(self._map) > len(other):
return False
if not isinstance(other, (set, frozenset)):
other = set(other)
return set(self._map.keys()) <= other
def __lt__(self, other):
if not isinstance(other, Iterable):
return False
if len(self._map) >= len(other):
return False
if not isinstance(other, (set, frozenset)):
other = set(other)
return set(self._map.keys()) < other
class StableSetEq(StableSet[T]):
"""
StableSetEq is a StableSet with a modified quality operator.
StableSetEq, like `set` and `dict_keys` [dict.keys()], and unlike OrderdSet,
disregards the items order when checking equality.
Unlike StableSet, `set`, or `dict_keys` - A StableSetEq can also equal be equal to a Sequence:
`StableSet([1, 2]) == [1, 2]` and `StableSet([1, 2]) == [2, 1]`; but `set([1, 2]) != [1, 2]`
"""
def __eq__(self, other: Any) -> bool:
"""
Returns true even if the containers don't have the same items in order.
Example:
>>> oset = StableSetEq([1, 3, 2])
>>> oset == [1, 3, 2]
True
>>> oset == [1, 2, 3]
True
>>> oset == [2, 3]
False
>>> oset == StableSetEq([3, 2, 1])
True
"""
if not isinstance(other, AbstractSet):
try:
other = set(other)
except TypeError:
# If `other` can't be converted into a set, it's not equal.
return False
return self._map.keys() == other
def __le__(self, other: SetLike[T]):
return len(self) <= len(other) and (
self._map.keys() <= other
if isinstance(other, AbstractSet)
else self._map.keys() <= set(other)
)
def __lt__(self, other: SetLike[T]):
return len(self) < len(other) and (
self._map.keys() < other
if isinstance(other, AbstractSet)
else self._map.keys() < set(other)
)
def __ge__(self, other: SetLike[T]):
return len(self) >= len(other) and (
self._map.keys() >= other
if isinstance(other, AbstractSet)
else self._map.keys() >= set(other)
)
def __gt__(self, other: SetLike[T]):
return len(self) > len(other) and (
self._map.keys() > other
if isinstance(other, AbstractSet)
else self._map.keys() > set(other)
)
class OrderedSet(StableSet[T]):
"""
An OrderedSet is a mutable data structure that is a hybrid of a list and a set.
It remembers its insertion order so that every entry has an index that can be looked up.
Featuring: O(1) Index lookup, insertion, iteration and membership testing.
But slow O(N) Deletion.
Using OrderedSet over StableSet is advised only if you require fast Index lookup -
Otherwise using StableSet is advised as it is much faster and has a smaller memory footprint.
In some aspects OrderedSet behaves like a `set` and in other aspects it behaves like a list.
Equality: OrderedSet, like `list` and `odict_keys` [OrderdDict.keys()], and unlike OrderdSet,
regards the items order when checking equality.
Unlike `set`, An OrderedSet can also equal be equal to a Sequence:
`StableSet([1, 2]) == [1, 2]` and `StableSet([1, 2]) != [2, 1]`; but `set([1, 2]) != [1, 2]`
The original implementation of OrderedSet was a recipe posted by Raymond Hettiger,
https://code.activestate.com/recipes/576694-orderedset/
Released under the MIT license.
Hettiger's implementation kept its content in a doubly-linked list referenced by a dict.
As a result, looking up an item by its index was an O(N) operation, while deletion was O(1).
This version makes different trade-offs for the sake of efficient lookups.
Its content is a standard Python list instead of a doubly-linked list.
This provides O(1) lookups by index at the expense of O(N) deletion,
as well as slightly faster iteration.
Example:
>>> OrderedSet([1, 1, 2, 3, 2])
OrderedSet([1, 2, 3])
"""
__slots__ = ("_items",)
_items: List[T]
def __init__(self, initial: Optional[SetInitializer[T]] = None):
self._items = []
self._map = {}
if initial is not None:
# In terms of duck-typing, the default __ior__ is compatible with
# the types we use, but it doesn't expect all the types we
# support as values for `initial`.
self |= initial # type: ignore
def __getitem__(self, index):
if isinstance(index, int):
return self._items[index]
elif isinstance(index, slice) and index == SLICE_ALL:
return self.copy()
elif isinstance(index, Iterable):
return [self._items[i] for i in index]
elif isinstance(index, slice) or hasattr(index, "__index__"):
result = self._items[index]
if isinstance(result, list):
return self.__class__(result)
else:
return result
else:
raise TypeError("Don't know how to index an OrderedSet by %r" % index)
def __eq__(self, other: Any) -> bool:
"""
Returns true if the containers have the same items.
If `other` is a Sequence, then order is checked, otherwise it is ignored.
Example:
>>> oset = OrderedSet([1, 3, 2])
>>> oset == [1, 3, 2]
True
>>> oset == [1, 2, 3]
False
>>> oset == [2, 3]
False
>>> oset == OrderedSet([3, 2, 1])
False
"""
if isinstance(other, Sequence):
# Check that this OrderedSet contains the same elements, in the
# same order, as the other object.
return len(self) == len(other) and self._items == list(other)
try:
other_as_set = set(other)
except TypeError:
# If `other` can't be converted into a set, it's not equal.
return False
else:
return self._map.keys() == other_as_set
def __le__(self, other: SetLike[T]):
return len(self) <= len(other) and (
self._map.keys() <= other
if isinstance(other, AbstractSet)
else self._items <= other
if isinstance(other, list)
else self._items <= list(other)
)
def __lt__(self, other: SetLike[T]):
return len(self) < len(other) and (
self._map.keys() < other
if isinstance(other, AbstractSet)
else self._items < other
if isinstance(other, list)
else self._items < list(other)
)
def __ge__(self, other: SetLike[T]):
return len(self) >= len(other) and (
self._map.keys() >= other
if isinstance(other, AbstractSet)
else self._items >= other
if isinstance(other, list)
else self._items >= list(other)
)
def __gt__(self, other: SetLike[T]):
return len(self) > len(other) and (
self._map.keys() > other
if isinstance(other, AbstractSet)
else self._items > other
if isinstance(other, list)
else self._items > list(other)
)
def clear(self) -> None:
del self._items[:]
self._map.clear()
def add(self, key: T) -> int:
if key not in self._map:
self._map[key] = len(self._items)
self._items.append(key)
return self._map[key]
def update(self, sequence: SetLike[T]) -> int:
item_index = 0
for item in sequence:
item_index = self.add(item)
return item_index
def index(self, key):
if isinstance(key, Iterable) and not _is_atomic(key):
return [self.index(subkey) for subkey in key]
return self._map[key]
def pop(self, index: int = -1) -> T:
if not self._items:
raise KeyError("Set is empty")
elem = self._items[index]
del self._items[index]
del self._map[elem]
return elem
def popitem(self, last: bool = True):
if not self._items:
raise KeyError("Set is empty")
index = -1 if last else 0
elem = self._items[index]
del self._items[index]
del self._map[elem]
return elem
def move_to_end(self, key):
if key in self:
self.discard(key)
self.add(key)
else:
raise KeyError(key)
def discard(self, key: T) -> None:
if key in self:
i = self._map[key]
del self._items[i]
del self._map[key]
for k, v in self._map.items():
if v >= i:
self._map[k] = v - 1
def _update_items(self, items: list) -> None:
"""
Replace the 'items' list of this OrderedSet with a new one, updating
self._map accordingly.
"""
self._items = items
self._map = {item: idx for (idx, item) in enumerate(items)}
def difference_update(self, *sets: SetLike[T]) -> None:
items_to_remove = set() # type: Set[T]
for other in sets:
items_as_set = set(other) # type: Set[T]
items_to_remove |= items_as_set
self._update_items(
[item for item in self._items if item not in items_to_remove]
)
def intersection_update(self, other: SetLike[T]) -> None:
other = set(other)
self._update_items([item for item in self._items if item in other])
def symmetric_difference_update(self, other: SetLike[T]) -> None:
items_to_add = [item for item in other if item not in self]
items_to_remove = set(other)
self._update_items(
[item for item in self._items if item not in items_to_remove] + items_to_add
)
class SortedSet:
def __init__(self, *args, set_=None, **kwargs):
self._sorted = None
if set_:
self.set_ = set_
else:
self.set_ = set(*args, **kwargs)
def __iter__(self):
yield from self._get_sorted()
def __str__(self) -> str:
if not self:
return f"{self.__class__.__name__}()"
return f"{self.__class__.__name__}({self._get_sorted()!r})"
__repr__ = __str__
def _get_sorted(self, reverse=False):
if self._sorted is None:
try:
self._sorted = sorted(self.set_, reverse=reverse)
except Exception:
self._sorted = sorted(self.set_, key=lambda x: str(x), reverse=reverse)
return self._sorted
def get(self):
"""
Get a random element
"""
return next(iter(self.set_))
def __and__(self, other):
# Intersection
result = self.set_ & other
return SortedSet(set_=result)
intersection = __and__
__rand__ = __and__
def intersection_update(self, other):
self.set_.intersection_update(other)
def __or__(self, other):
# Union
result = self.set_ | other
return SortedSet(set_=result)
union = __ror__ = __or__
def __sub__(self, other):
# Difference
result = self.set_ - other
return SortedSet(set_=result)
difference = __sub__
def difference_update(self, *sets):
self._sorted = None
self.set_.difference_update(*sets)
def isdisjoint(self, other):
return self.set_.isdisjoint(other)
def __xor__(self, other):
# Symmetric difference
result = self.set_ ^ other
return SortedSet(set_=result)
symmetric_difference = __rxor__ = __xor__
def symmetric_difference_update(self, other):
self._sorted = None
self.set_.symmetric_difference_update(other)
def __rsub__(self, other):
result = other - self.set_
return SortedSet(set_=result)
def add(self, item):
self._sorted = None
self.set_.add(item)
def clear(self):
self._sorted = None
self.set_.clear()
def discard(self, key):
self._sorted = None
self.set_.discard(key)
def copy(self):
return SortedSet(set_=set(self.set_))
def __le__(self, other):
return self.set_.issubset(other)
issubset = __le__
def __lt__(self, other):
return self.set_ < other
def __ge__(self, other):
return self.set_.issuperset(other)
issuperset = __ge__
def __gt__(self, other):
return self.set_ > other
def remove(self, key):
self._sorted = None
self.set_.remove(key)
def update(self, sequence):
self._sorted = None
self.set_.update(sequence)
def __len__(self):
return len(self.set_) if self.set_ else 0
def __eq__(self, other):
if not isinstance(other, Iterable):
return False
if len(self.set_) != len(other):
return False
if isinstance(other, SortedSet):
return self.set_ == other.set_
if not isinstance(other, (set, frozenset)):
other = set(other)
return self.set_ == other
def __reversed__(self):
if self._sorted is None:
self._get_sorted()
return reversed(self._sorted)
def __getitem__(self, index):
items = self._get_sorted()
if isinstance(index, int):
return items[index]
elif isinstance(index, slice) and index == SLICE_ALL:
return self.copy()
if isinstance(index, Iterable):
return [items[i] for i in index]
elif isinstance(index, slice) or hasattr(index, "__index__"):
result = items[index]
if isinstance(result, list):
return self.__class__(result)
else:
return result
else:
raise TypeError(f"Don't know how to index a SortedSet by {index}")
def index(self, key): # NOQA
"""
Get the index of a given entry, raising an IndexError if it's not present
`key` can be an iterable of entries that is not a string, in which case
this returns a list of indices.
Example:
>>> oset = StableSet([1, 2, 3])
>>> oset.index(2)
1
"""
items = self._get_sorted()
try:
if isinstance(key, Iterable) and not _is_atomic(key):
return [self.index(subkey) for subkey in key]
for index, item in enumerate(items):
if item == key:
return index
raise KeyError(key)
# return list(self._map.keys()).index(key)
except ValueError:
raise KeyError(key)
# Provide some compatibility with pd.Index
get_loc = index
get_indexer = index
def pop(self, index=None):
if index is None:
self._sorted = None
return self.set_.pop()
items = self._get_sorted()
result = items.pop(index)
self.set_.remove(result)
return result
def isorderedsubset(self: SetLike, other: SetLike, non_consecutive: bool = False):
if len(self) > len(other):
return False
if non_consecutive:
i = 0
self_len = len(self)
for other_item in other:
if other_item == self[i]:
i += 1
if i == self_len:
return True
return False
else:
for self_item, other_item in zip(self, other):
if not self_item == other_item:
return False
return True
def isorderedsuperset(self, other: SetLike, non_consecutive: bool = False):
return StableSet.isorderedsubset(other, self, non_consecutive)
