You can not select more than 25 topics
Topics must start with a letter or number, can include dashes ('-') and can be up to 35 characters long.
257 lines
9.2 KiB
257 lines
9.2 KiB
"""
|
|
TwoWayMap implements mapping from keys to values, and values back to keys. Since keys and values
|
|
are not really different here, they are referred to throughout as 'left' and 'right' values.
|
|
|
|
TwoWayMap supports different types of containers when one value maps to multiple. You may add
|
|
support for additional container types using register_container() module function.
|
|
|
|
It's implemented using Python dictionaries, so both 'left' and 'right' values must be hashable.
|
|
|
|
For example, to create a dictionary-like structure mapping one key to one value, and which allows
|
|
to quickly tell the set of keys that map to a given value, we can use m=TwoWayMap(left=set,
|
|
right="single"). Then m.insert(key, value) sets the given key to the given value (overwriting the
|
|
value previously set, since the "right" dataset is "single" values), m.lookup_left(key) returns
|
|
that value, and m.lookup_right(value) returns a `set` of keys that map to the value.
|
|
"""
|
|
|
|
import six
|
|
|
|
# Special sentinel value which can never be legitimately stored in TwoWayMap, to easily tell the
|
|
# difference between a present and absent value.
|
|
_NIL = object()
|
|
|
|
|
|
class TwoWayMap(object):
|
|
def __init__(self, left=set, right=set):
|
|
"""
|
|
Create a new TwoWayMap. The `left` and `right` parameters determine the type of bin for
|
|
storing multiple values on the respective side of the map. E.g. if right=set, then
|
|
lookup_left() will return a set (what's on the right side). Supported values are:
|
|
set: a set of values.
|
|
list: a list of values, with new items added at the end of the list.
|
|
"single": a single value, new items overwrite previous ones.
|
|
"strict": a single value, new items must not overwrite previous ones.
|
|
|
|
To add support for another bin type, use twowaymap.register_container().
|
|
|
|
E.g. for TwoWayMap(left="single", right="strict"),
|
|
after insert(1, "a"), insert(1, "b") will succeed, but insert(2, "a") will fail.
|
|
|
|
E.g. for TwoWayMap(left=list, right="single"),
|
|
after insert(1, "a"), insert(1, "b"), insert(2, "a"),
|
|
lookup_left(1) will return ["a", "b"], and lookup_right("a") will return 2.
|
|
"""
|
|
self._left_bin = _mapper_types[left]
|
|
self._right_bin = _mapper_types[right]
|
|
self._fwd = {}
|
|
self._bwd = {}
|
|
|
|
def __nonzero__(self):
|
|
return bool(self._fwd)
|
|
|
|
__bool__ = __nonzero__
|
|
|
|
def lookup_left(self, left, default=None):
|
|
""" Returns the value(s) on the right corresponding to the given value on the left. """
|
|
return self._fwd.get(left, default)
|
|
|
|
def lookup_right(self, right, default=None):
|
|
""" Returns the value(s) on the left corresponding to the given value on the right. """
|
|
return self._bwd.get(right, default)
|
|
|
|
def count_left(self):
|
|
""" Returns the count of unique values on the left."""
|
|
return len(self._fwd)
|
|
|
|
def count_right(self):
|
|
""" Returns the count of unique values on the right."""
|
|
return len(self._bwd)
|
|
|
|
def left_all(self):
|
|
""" Returns an iterable over all values on the left."""
|
|
return six.iterkeys(self._fwd)
|
|
|
|
def right_all(self):
|
|
""" Returns an iterable over all values on the right."""
|
|
return six.iterkeys(self._bwd)
|
|
|
|
def insert(self, left, right):
|
|
""" Insert the (left, right) value pair. """
|
|
# The tricky thing here is to keep the two maps consistent if an update to the second one
|
|
# raises an exception. To handle it, add_item must return what got added and removed, so that
|
|
# we can restore things after an exception. An exception could be caused by a "strict" bin
|
|
# type, or by using an un-hashable key (on either left or right side), or by using a custom
|
|
# container that can throw.
|
|
right_removed, right_added = self._right_bin.add_item(self._fwd, left, right)
|
|
try:
|
|
left_removed, _ = self._left_bin.add_item(self._bwd, right, left)
|
|
except:
|
|
# _left_bin is responsible to stay unchanged if there was an exception. Now we need to bring
|
|
# _right_bin back in sync with _left_bin.
|
|
if right_added is not _NIL:
|
|
self._right_bin.remove_item(self._fwd, left, right_added)
|
|
if right_removed is not _NIL:
|
|
self._right_bin.add_item(self._fwd, left, right_removed)
|
|
raise
|
|
|
|
# It's possible for add_item to overwrite elements, in which case we need to remove the
|
|
# other side of the mapping for the removed element.
|
|
if right_removed is not _NIL:
|
|
self._left_bin.remove_item(self._bwd, right_removed, left)
|
|
if left_removed is not _NIL:
|
|
self._right_bin.remove_item(self._fwd, left_removed, right)
|
|
|
|
def remove(self, left, right):
|
|
""" Remove the (left, right) value pair. """
|
|
self._right_bin.remove_item(self._fwd, left, right)
|
|
self._left_bin.remove_item(self._bwd, right, left)
|
|
|
|
def remove_left(self, left):
|
|
""" Remove all values on the right corresponding to the given value on the left. """
|
|
right_removed = self._right_bin.remove_key(self._fwd, left)
|
|
for x in right_removed:
|
|
self._left_bin.remove_item(self._bwd, x, left)
|
|
|
|
def remove_right(self, right):
|
|
""" Remove all values on the left corresponding to the given value on the right. """
|
|
left_removed = self._left_bin.remove_key(self._bwd, right)
|
|
for x in left_removed:
|
|
self._right_bin.remove_item(self._fwd, x, right)
|
|
|
|
def clear(self):
|
|
""" Clear the entire map. """
|
|
self._fwd.clear()
|
|
self._bwd.clear()
|
|
|
|
#----------------------------------------------------------------------
|
|
# The private classes below implement the different container types.
|
|
|
|
class _BaseBinType(object):
|
|
""" Base class for other BinTypes. """
|
|
def add_item(self, mapping, key, value):
|
|
pass
|
|
def remove_item(self, mapping, key, value):
|
|
pass
|
|
def remove_key(self, mapping, key):
|
|
pass
|
|
|
|
|
|
class _SingleValueBin(_BaseBinType):
|
|
""" Bin that contains a single value, with new values overwriting previous ones."""
|
|
def add_item(self, mapping, key, value):
|
|
stored = mapping.get(key, _NIL)
|
|
mapping[key] = value
|
|
if stored is _NIL:
|
|
return _NIL, value
|
|
elif stored == value:
|
|
return _NIL, _NIL
|
|
else:
|
|
return stored, value
|
|
|
|
def remove_item(self, mapping, key, value):
|
|
stored = mapping.get(key, _NIL)
|
|
if stored == value:
|
|
del mapping[key]
|
|
|
|
def remove_key(self, mapping, key):
|
|
stored = mapping.pop(key, _NIL)
|
|
return () if stored is _NIL else (stored,)
|
|
|
|
|
|
class _SingleValueStrictBin(_SingleValueBin):
|
|
""" Bin that contains a single value, overwriting which raises ValueError."""
|
|
def add_item(self, mapping, key, value):
|
|
stored = mapping.get(key, _NIL)
|
|
if stored is _NIL:
|
|
mapping[key] = value
|
|
return _NIL, value
|
|
elif stored == value:
|
|
return _NIL, _NIL
|
|
else:
|
|
raise ValueError("twowaymap: one-to-one map violation for key %s" % key)
|
|
|
|
|
|
class _ContainerBin(_BaseBinType):
|
|
"""
|
|
Bin that contains a container of values managed by the passed-in functions. See
|
|
register_container() for documentation of the arguments.
|
|
"""
|
|
def __init__(self, make_func, add_func, remove_func):
|
|
self.make = make_func
|
|
self.add = add_func
|
|
self.remove = remove_func
|
|
|
|
def add_item(self, mapping, key, value):
|
|
stored = mapping.get(key, _NIL)
|
|
if stored is _NIL:
|
|
mapping[key] = self.make(value)
|
|
return _NIL, value
|
|
else:
|
|
return _NIL, (value if self.add(stored, value) else _NIL)
|
|
|
|
def remove_item(self, mapping, key, value):
|
|
stored = mapping.get(key, _NIL)
|
|
if stored is not _NIL:
|
|
self.remove(stored, value)
|
|
if not stored:
|
|
del mapping[key]
|
|
|
|
def remove_key(self, mapping, key):
|
|
return mapping.pop(key, ())
|
|
|
|
#----------------------------------------------------------------------
|
|
|
|
_mapper_types = {
|
|
'single': _SingleValueBin(),
|
|
'strict': _SingleValueStrictBin(),
|
|
}
|
|
|
|
def register_container(cls, make_func, add_func, remove_func):
|
|
"""
|
|
Register another container type. The first argument can be the container's class object, but
|
|
really can be any hashable value, which you can then give as an argument to left= or right=
|
|
arguments when constructing a TwoWayMap. The other arguments are:
|
|
|
|
make_func(value) - must return a new instance of the container with a single value.
|
|
This container must support iteration through values, and in boolean context must
|
|
evaluate to whether it's non-empty.
|
|
|
|
add_func(container, value) - must add value to container, only if it's not already there,
|
|
and return True if the value was added, False if it was already there.
|
|
|
|
remove_func(container, value) - must remove value from container if present.
|
|
This must never raise an exception, since that could leave the map in inconsistent state.
|
|
"""
|
|
_mapper_types[cls] = _ContainerBin(make_func, add_func, remove_func)
|
|
|
|
|
|
# Allow `set` to be used as a bin type.
|
|
def _set_make(value):
|
|
return {value}
|
|
def _set_add(container, value):
|
|
if value not in container:
|
|
container.add(value)
|
|
return True
|
|
return False
|
|
def _set_remove(container, value):
|
|
container.discard(value)
|
|
|
|
register_container(set, _set_make, _set_add, _set_remove)
|
|
|
|
|
|
# Allow `list` to be used as a bin type.
|
|
def _list_make(value):
|
|
return [value]
|
|
def _list_add(container, value):
|
|
if value not in container:
|
|
container.append(value)
|
|
return True
|
|
return False
|
|
def _list_remove(container, value):
|
|
try:
|
|
container.remove(value)
|
|
except ValueError:
|
|
pass
|
|
|
|
register_container(list, _list_make, _list_add, _list_remove)
|