import logging
import json
from collections import defaultdict
from typing import DefaultDict, Generic, TypeVar, List
K = TypeVar("K")
V = TypeVar("V")
logger = logging.getLogger(__name__)
def _str_dump_multimap(mm): # pragma: no cover
d = defaultdict(list)
for k, v in mm.items():
d[k].append(repr(v))
return json.dumps(d, sort_keys=True, indent=4)
[docs]
class MultiMap(Generic[K, V]):
"""Implements a simple MultiMap data structure on top of a dictionary of lists"""
__slots__ = ['contents', 'clean']
contents: DefaultDict[K, List[V]]
clean: bool
def __init__(self, **kwargs):
self.contents = defaultdict(list)
for k, v in kwargs.items():
self.contents[k].append(v)
self.invalidate()
def invalidate(self):
self.clean = False
def __getitem__(self, key):
return self.contents[key]
def __setitem__(self, key, value):
self.contents[key].append(value)
self.invalidate()
def pop(self, key, value):
"""
Removes `value` from the collection of values stored at `key`
and returns the `tuple` `(key, value)`
Raises
------
IndexError
KeyError
"""
objs = self.contents[key]
objs.pop(objs.index(value))
if len(objs) == 0:
self.contents.pop(key)
self.invalidate()
return len(objs)
def popv(self, value):
for k in self:
if value in self[k]:
self.pop(k, value)
return len(self[k])
self.invalidate()
return None
def __iter__(self):
return iter(self.contents)
def __contains__(self, key):
return key in self.contents
def keys(self):
"""
Returns an iterator over the keys of :attr:`contents`
An alias of :meth:`__iter__`
"""
return iter(self)
def values(self):
"""
Returns an iterator over the values of :attr:`contents`
"""
for k in self:
for v in self[k]:
yield v
def items(self):
"""
Returns an iterator over the items of :attr:`contents`. Each item
takes the form of `(key, value)`.
"""
for k in self:
for v in self[k]:
yield (k, v)
def lists(self):
return self.contents.items()
def __len__(self):
"""
Returns the number of items in :attr:`contents`
"""
return sum(len(self[k]) for k in self)
def __repr__(self): # pragma: no cover
return _str_dump_multimap(self)
def __eq__(self, other):
if other is None: # pragma: no cover
return False
for a in self.keys():
if a in other:
if self[a] != other[a]:
return False
else:
if self[a] != []:
return False
for b in other.keys():
if b in self:
continue
else:
if other[b] != []:
return False
# for a, b in izip_longest(self.items(), other.items()):
# if a != b:
# return False
return True
def __ne__(self, other):
return not self == other
def update(self, mapping):
for k, v in mapping.items():
self[k] = v
self.invalidate()
def has_value(self, value):
for v in self.values():
if v == value:
return True
return False
def __reduce__(self):
return self.__class__, (), self.__getstate__()
def __getstate__(self):
return self.contents
def __setstate__(self, state):
self.contents = state
self.invalidate()
def copy(self):
new = self.__class__()
for k, v in self.items():
new[k].extend(v)
return new
def clone(self):
return self.copy()
[docs]
class OrderedMultiMap(MultiMap[K, V]):
"""
Implements a simple MultiMap data structure on top of a dictionary of lists
that remembers the order keys were first inserted in.
"""
__slots__ = ['key_order', ]
key_order: List[K]
def __init__(self, **kwargs):
self.contents = defaultdict(list)
self.key_order = []
for k, v in kwargs.items():
if k not in self.key_order:
self.key_order.append(k)
self.contents[k].append(v)
self.invalidate()
def __iter__(self):
"""
Returns an iterator over the keys of :attr:`contents` in the order
they were added.
"""
for key in self.key_order:
yield key
#: Alias of :meth:`__iter__`
keys = __iter__
def values(self):
for key in self.key_order:
for v in self[key]:
yield v
def items(self):
"""
As in :class:`MultiMap`, but items are yielded in the order their keys were
first added.
"""
for key in self.key_order:
for v in self[key]:
yield (key, v)
def lists(self):
for key in self.key_order:
yield key, self[key]
def reorder(self, new_order):
assert set(new_order) == set(self.key_order)
self.key_order = list(new_order)
def __setitem__(self, key, value):
if key not in self.key_order:
self.key_order.append(key)
self.contents[key].append(value)
self.invalidate()
def __repr__(self): # pragma: no cover
return ''.join((repr(self.key_order), '\n', _str_dump_multimap(self)))
def __getstate__(self):
return self.contents, self.key_order
def __setstate__(self, state):
self.contents, self.key_order = state
self.invalidate()