123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199 |
- """
- An OrderedSet is a custom MutableSet that remembers its order, so that every
- entry has an index that can be looked up.
- Based on a recipe originally posted to ActiveState Recipes by Raymond Hettiger,
- and released under the MIT license.
- Rob Speer's changes are as follows:
- - changed the content from a doubly-linked list to a regular Python list.
- Seriously, who wants O(1) deletes but O(N) lookups by index?
- - add() returns the index of the added item
- - index() just returns the index of an item
- - added a __getstate__ and __setstate__ so it can be pickled
- - added __getitem__
- """
- import collections
- SLICE_ALL = slice(None)
- __version__ = '2.0.1'
- def is_iterable(obj):
- """
- Are we being asked to look up a list of things, instead of a single thing?
- We check for the `__iter__` attribute so that this can cover types that
- don't have to be known by this module, such as NumPy arrays.
- Strings, however, should be considered as atomic values to look up, not
- iterables. The same goes for tuples, since they are immutable and therefore
- valid entries.
- We don't need to check for the Python 2 `unicode` type, because it doesn't
- have an `__iter__` attribute anyway.
- """
- return hasattr(obj, '__iter__') and not isinstance(obj, str) and not isinstance(obj, tuple)
- class OrderedSet(collections.MutableSet):
- """
- An OrderedSet is a custom MutableSet that remembers its order, so that
- every entry has an index that can be looked up.
- """
- def __init__(self, iterable=None):
- self.items = []
- self.map = {}
- if iterable is not None:
- self |= iterable
- def __len__(self):
- return len(self.items)
- def __getitem__(self, index):
- """
- Get the item at a given index.
- If `index` is a slice, you will get back that slice of items. If it's
- the slice [:], exactly the same object is returned. (If you want an
- independent copy of an OrderedSet, use `OrderedSet.copy()`.)
- If `index` is an iterable, you'll get the OrderedSet of items
- corresponding to those indices. This is similar to NumPy's
- "fancy indexing".
- """
- if index == SLICE_ALL:
- return self
- elif hasattr(index, '__index__') or isinstance(index, slice):
- result = self.items[index]
- if isinstance(result, list):
- return OrderedSet(result)
- else:
- return result
- elif is_iterable(index):
- return OrderedSet([self.items[i] for i in index])
- else:
- raise TypeError("Don't know how to index an OrderedSet by %r" %
- index)
- def copy(self):
- return OrderedSet(self)
- def __getstate__(self):
- if len(self) == 0:
- # 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 OrderedSet.
- return (None,)
- else:
- return list(self)
- def __setstate__(self, state):
- if state == (None,):
- self.__init__([])
- else:
- self.__init__(state)
- def __contains__(self, key):
- return key in self.map
- def add(self, key):
- """
- Add `key` as an item to this OrderedSet, then return its index.
- If `key` is already in the OrderedSet, return the index it already
- had.
- """
- if key not in self.map:
- self.map[key] = len(self.items)
- self.items.append(key)
- return self.map[key]
- append = add
- def update(self, sequence):
- """
- Update the set with the given iterable sequence, then return the index
- of the last element inserted.
- """
- item_index = None
- try:
- for item in sequence:
- item_index = self.add(item)
- except TypeError:
- raise ValueError('Argument needs to be an iterable, got %s' % type(sequence))
- return item_index
- def index(self, key):
- """
- 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.
- """
- if is_iterable(key):
- return [self.index(subkey) for subkey in key]
- return self.map[key]
- def pop(self):
- """
- Remove and return the last element from the set.
- Raises KeyError if the set is empty.
- """
- if not self.items:
- raise KeyError('Set is empty')
- elem = self.items[-1]
- del self.items[-1]
- del self.map[elem]
- return elem
- def discard(self, key):
- """
- 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.
- """
- if key in self:
- i = self.items.index(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 clear(self):
- """
- Remove all items from this OrderedSet.
- """
- del self.items[:]
- self.map.clear()
- def __iter__(self):
- return iter(self.items)
- def __reversed__(self):
- return reversed(self.items)
- def __repr__(self):
- if not self:
- return '%s()' % (self.__class__.__name__,)
- return '%s(%r)' % (self.__class__.__name__, list(self))
- def __eq__(self, other):
- if isinstance(other, OrderedSet):
- return len(self) == len(other) and self.items == other.items
- 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 set(self) == other_as_set
|