tahoe-lafs/src/allmydata/util/dictutil.py
2010-02-26 01:14:33 -07:00

648 lines
21 KiB
Python

"""
Tools to mess with dicts.
"""
import copy, operator
from bisect import bisect_left, insort_left
from allmydata.util.assertutil import _assert, precondition
def move(k, d1, d2, strict=False):
"""
Move item with key k from d1 to d2.
"""
if strict and not d1.has_key(k):
raise KeyError, k
d2[k] = d1[k]
del d1[k]
def subtract(d1, d2):
"""
Remove all items from d1 whose key occurs in d2.
@returns d1
"""
if len(d1) > len(d2):
for k in d2.keys():
if d1.has_key(k):
del d1[k]
else:
for k in d1.keys():
if d2.has_key(k):
del d1[k]
return d1
class DictOfSets(dict):
def add(self, key, value):
if key in self:
self[key].add(value)
else:
self[key] = set([value])
def union(self, key, values):
if key in self:
self[key].update(values)
else:
self[key] = set(values)
def update(self, otherdictofsets):
for key, value in otherdictofsets.iteritems():
self.union(key, value)
def discard(self, key, value):
if not key in self:
return
self[key].discard(value)
if not self[key]:
del self[key]
class UtilDict:
def __init__(self, initialdata={}):
self.d = {}
self.update(initialdata)
def del_if_present(self, key):
if self.has_key(key):
del self[key]
def items_sorted_by_value(self):
"""
@return a sequence of (key, value,) pairs sorted according to value
"""
l = [(x[1], x[0],) for x in self.d.iteritems()]
l.sort()
return [(x[1], x[0],) for x in l]
def items_sorted_by_key(self):
"""
@return a sequence of (key, value,) pairs sorted according to key
"""
l = self.d.items()
l.sort()
return l
def __repr__(self, *args, **kwargs):
return self.d.__repr__(*args, **kwargs)
def __str__(self, *args, **kwargs):
return self.d.__str__(*args, **kwargs)
def __contains__(self, *args, **kwargs):
return self.d.__contains__(*args, **kwargs)
def __len__(self, *args, **kwargs):
return self.d.__len__(*args, **kwargs)
def __cmp__(self, other):
try:
return self.d.__cmp__(other)
except TypeError, le:
# maybe we should look for a .d member in other. I know this is insanely kludgey, but the Right Way To Do It is for dict.__cmp__ to use structural typing ("duck typing")
try:
return self.d.__cmp__(other.d)
except:
raise le
def __eq__(self, *args, **kwargs):
return self.d.__eq__(*args, **kwargs)
def __ne__(self, *args, **kwargs):
return self.d.__ne__(*args, **kwargs)
def __gt__(self, *args, **kwargs):
return self.d.__gt__(*args, **kwargs)
def __ge__(self, *args, **kwargs):
return self.d.__ge__(*args, **kwargs)
def __le__(self, *args, **kwargs):
return self.d.__le__(*args, **kwargs)
def __lt__(self, *args, **kwargs):
return self.d.__lt__(*args, **kwargs)
def __getitem__(self, *args, **kwargs):
return self.d.__getitem__(*args, **kwargs)
def __setitem__(self, *args, **kwargs):
return self.d.__setitem__(*args, **kwargs)
def __delitem__(self, *args, **kwargs):
return self.d.__delitem__(*args, **kwargs)
def __iter__(self, *args, **kwargs):
return self.d.__iter__(*args, **kwargs)
def clear(self, *args, **kwargs):
return self.d.clear(*args, **kwargs)
def copy(self, *args, **kwargs):
return self.__class__(self.d.copy(*args, **kwargs))
def fromkeys(self, *args, **kwargs):
return self.__class__(self.d.fromkeys(*args, **kwargs))
def get(self, key, default=None):
return self.d.get(key, default)
def has_key(self, *args, **kwargs):
return self.d.has_key(*args, **kwargs)
def items(self, *args, **kwargs):
return self.d.items(*args, **kwargs)
def iteritems(self, *args, **kwargs):
return self.d.iteritems(*args, **kwargs)
def iterkeys(self, *args, **kwargs):
return self.d.iterkeys(*args, **kwargs)
def itervalues(self, *args, **kwargs):
return self.d.itervalues(*args, **kwargs)
def keys(self, *args, **kwargs):
return self.d.keys(*args, **kwargs)
def pop(self, *args, **kwargs):
return self.d.pop(*args, **kwargs)
def popitem(self, *args, **kwargs):
return self.d.popitem(*args, **kwargs)
def setdefault(self, *args, **kwargs):
return self.d.setdefault(*args, **kwargs)
def update(self, *args, **kwargs):
self.d.update(*args, **kwargs)
def values(self, *args, **kwargs):
return self.d.values(*args, **kwargs)
class NumDict:
def __init__(self, initialdict={}):
self.d = copy.deepcopy(initialdict)
def add_num(self, key, val, default=0):
"""
If the key doesn't appear in self then it is created with value default
(before addition).
"""
self.d[key] = self.d.get(key, default) + val
def subtract_num(self, key, val, default=0):
self.d[key] = self.d.get(key, default) - val
def sum(self):
"""
@return: the sum of all values
"""
return reduce(operator.__add__, self.d.values())
def inc(self, key, default=0):
"""
Increment the value associated with key in dict. If there is no such
key, then one will be created with initial value 0 (before inc() --
therefore value 1 after inc).
"""
self.add_num(key, 1, default)
def dec(self, key, default=0):
"""
Decrement the value associated with key in dict. If there is no such
key, then one will be created with initial value 0 (before dec() --
therefore value -1 after dec).
"""
self.subtract_num(key, 1, default)
def items_sorted_by_value(self):
"""
@return a sequence of (key, value,) pairs sorted according to value
"""
l = [(x[1], x[0],) for x in self.d.iteritems()]
l.sort()
return [(x[1], x[0],) for x in l]
def item_with_largest_value(self):
it = self.d.iteritems()
(winner, winnerval,) = it.next()
try:
while True:
n, nv = it.next()
if nv > winnerval:
winner = n
winnerval = nv
except StopIteration:
pass
return (winner, winnerval,)
def items_sorted_by_key(self):
"""
@return a sequence of (key, value,) pairs sorted according to key
"""
l = self.d.items()
l.sort()
return l
def __repr__(self, *args, **kwargs):
return self.d.__repr__(*args, **kwargs)
def __str__(self, *args, **kwargs):
return self.d.__str__(*args, **kwargs)
def __contains__(self, *args, **kwargs):
return self.d.__contains__(*args, **kwargs)
def __len__(self, *args, **kwargs):
return self.d.__len__(*args, **kwargs)
def __cmp__(self, other):
try:
return self.d.__cmp__(other)
except TypeError, le:
# maybe we should look for a .d member in other. I know this is insanely kludgey, but the Right Way To Do It is for dict.__cmp__ to use structural typing ("duck typing")
try:
return self.d.__cmp__(other.d)
except:
raise le
def __eq__(self, *args, **kwargs):
return self.d.__eq__(*args, **kwargs)
def __ne__(self, *args, **kwargs):
return self.d.__ne__(*args, **kwargs)
def __gt__(self, *args, **kwargs):
return self.d.__gt__(*args, **kwargs)
def __ge__(self, *args, **kwargs):
return self.d.__ge__(*args, **kwargs)
def __le__(self, *args, **kwargs):
return self.d.__le__(*args, **kwargs)
def __lt__(self, *args, **kwargs):
return self.d.__lt__(*args, **kwargs)
def __getitem__(self, *args, **kwargs):
return self.d.__getitem__(*args, **kwargs)
def __setitem__(self, *args, **kwargs):
return self.d.__setitem__(*args, **kwargs)
def __delitem__(self, *args, **kwargs):
return self.d.__delitem__(*args, **kwargs)
def __iter__(self, *args, **kwargs):
return self.d.__iter__(*args, **kwargs)
def clear(self, *args, **kwargs):
return self.d.clear(*args, **kwargs)
def copy(self, *args, **kwargs):
return self.__class__(self.d.copy(*args, **kwargs))
def fromkeys(self, *args, **kwargs):
return self.__class__(self.d.fromkeys(*args, **kwargs))
def get(self, key, default=0):
return self.d.get(key, default)
def has_key(self, *args, **kwargs):
return self.d.has_key(*args, **kwargs)
def items(self, *args, **kwargs):
return self.d.items(*args, **kwargs)
def iteritems(self, *args, **kwargs):
return self.d.iteritems(*args, **kwargs)
def iterkeys(self, *args, **kwargs):
return self.d.iterkeys(*args, **kwargs)
def itervalues(self, *args, **kwargs):
return self.d.itervalues(*args, **kwargs)
def keys(self, *args, **kwargs):
return self.d.keys(*args, **kwargs)
def pop(self, *args, **kwargs):
return self.d.pop(*args, **kwargs)
def popitem(self, *args, **kwargs):
return self.d.popitem(*args, **kwargs)
def setdefault(self, *args, **kwargs):
return self.d.setdefault(*args, **kwargs)
def update(self, *args, **kwargs):
return self.d.update(*args, **kwargs)
def values(self, *args, **kwargs):
return self.d.values(*args, **kwargs)
def del_if_present(d, k):
if d.has_key(k):
del d[k]
class ValueOrderedDict:
"""
Note: this implementation assumes that the values do not mutate and change
their sort order. That is, it stores the values in a sorted list and
as items are added and removed from the dict, it makes updates to the list
which will keep the list sorted. But if a value that is currently sitting
in the list changes its sort order, then the internal consistency of this
object will be lost.
If that happens, and if assertion checking is turned on, then you will get
an assertion failure the very next time you try to do anything with this
ValueOrderedDict. However, those internal consistency checks are very slow
and almost certainly unacceptable to leave turned on in production code.
"""
class ItemIterator:
def __init__(self, c):
self.c = c
self.i = 0
def __iter__(self):
return self
def next(self):
precondition(self.i <= len(self.c.l), "The iterated ValueOrderedDict doesn't have this many elements. Most likely this is because someone altered the contents of the ValueOrderedDict while the iteration was in progress.", self.i, self.c)
precondition((self.i == len(self.c.l)) or self.c.d.has_key(self.c.l[self.i][1]), "The iterated ValueOrderedDict doesn't have this key. Most likely this is because someone altered the contents of the ValueOrderedDict while the iteration was in progress.", self.i, (self.i < len(self.c.l)) and self.c.l[self.i], self.c)
if self.i == len(self.c.l):
raise StopIteration
le = self.c.l[self.i]
self.i += 1
return (le[1], le[0],)
def iteritems(self):
return ValueOrderedDict.ItemIterator(self)
def items(self):
return zip(map(operator.__getitem__, self.l, [1]*len(self.l)), map(operator.__getitem__, self.l, [0]*len(self.l)))
def values(self):
return map(operator.__getitem__, self.l, [0]*len(self.l))
def keys(self):
return map(operator.__getitem__, self.l, [1]*len(self.l))
class KeyIterator:
def __init__(self, c):
self.c = c
self.i = 0
def __iter__(self):
return self
def next(self):
precondition(self.i <= len(self.c.l), "The iterated ValueOrderedDict doesn't have this many elements. Most likely this is because someone altered the contents of the ValueOrderedDict while the iteration was in progress.", self.i, self.c)
precondition((self.i == len(self.c.l)) or self.c.d.has_key(self.c.l[self.i][1]), "The iterated ValueOrderedDict doesn't have this key. Most likely this is because someone altered the contents of the ValueOrderedDict while the iteration was in progress.", self.i, (self.i < len(self.c.l)) and self.c.l[self.i], self.c)
if self.i == len(self.c.l):
raise StopIteration
le = self.c.l[self.i]
self.i += 1
return le[1]
def iterkeys(self):
return ValueOrderedDict.KeyIterator(self)
class ValueIterator:
def __init__(self, c):
self.c = c
self.i = 0
def __iter__(self):
return self
def next(self):
precondition(self.i <= len(self.c.l), "The iterated ValueOrderedDict doesn't have this many elements. Most likely this is because someone altered the contents of the ValueOrderedDict while the iteration was in progress.", self.i, self.c)
precondition((self.i == len(self.c.l)) or self.c.d.has_key(self.c.l[self.i][1]), "The iterated ValueOrderedDict doesn't have this key. Most likely this is because someone altered the contents of the ValueOrderedDict while the iteration was in progress.", self.i, (self.i < len(self.c.l)) and self.c.l[self.i], self.c)
if self.i == len(self.c.l):
raise StopIteration
le = self.c.l[self.i]
self.i += 1
return le[0]
def itervalues(self):
return ValueOrderedDict.ValueIterator(self)
def __init__(self, initialdata={}):
self.d = {} # k: key, v: val
self.l = [] # sorted list of tuples of (val, key,)
self.update(initialdata)
assert self._assert_invariants()
def __len__(self):
return len(self.l)
def __repr_n__(self, n=None):
s = ["{",]
try:
iter = self.iteritems()
x = iter.next()
s.append(str(x[0])); s.append(": "); s.append(str(x[1]))
i = 1
while (n is None) or (i < n):
x = iter.next()
s.append(", "); s.append(str(x[0])); s.append(": "); s.append(str(x[1]))
except StopIteration:
pass
s.append("}")
return ''.join(s)
def __repr__(self):
return "<%s %s>" % (self.__class__.__name__, self.__repr_n__(),)
def __str__(self):
return "<%s %s>" % (self.__class__.__name__, self.__repr_n__(16),)
def __eq__(self, other):
for (k, v,) in other.iteritems():
if not self.d.has_key(k) or self.d[k] != v:
return False
return True
def __ne__(self, other):
return not self.__eq__(other)
def _assert_invariants(self):
iter = self.l.__iter__()
try:
oldx = iter.next()
while True:
x = iter.next()
# self.l is required to be sorted
_assert(x >= oldx, x, oldx)
# every element of self.l is required to appear in self.d
_assert(self.d.has_key(x[1]), x)
oldx =x
except StopIteration:
pass
for (k, v,) in self.d.iteritems():
i = bisect_left(self.l, (v, k,))
while (self.l[i][0] is not v) or (self.l[i][1] is not k):
i += 1
_assert(i < len(self.l), i, len(self.l), k, v, self.l)
_assert(self.l[i][0] is v, i, v, l=self.l, d=self.d)
_assert(self.l[i][1] is k, i, k, l=self.l, d=self.d)
return True
def insert(self, key, val=None):
assert self._assert_invariants()
result = self.__setitem__(key, val)
assert self._assert_invariants()
return result
def setdefault(self, key, default=None):
assert self._assert_invariants()
if not self.has_key(key):
self[key] = default
assert self._assert_invariants()
return self[key]
def __setitem__(self, key, val=None):
assert self._assert_invariants()
if self.d.has_key(key):
oldval = self.d[key]
if oldval != val:
# re-sort
i = bisect_left(self.l, (oldval, key,))
while (self.l[i][0] is not oldval) or (self.l[i][1] is not key):
i += 1
self.l.pop(i)
insort_left(self.l, (val, key,))
elif oldval is not val:
# replace
i = bisect_left(self.l, (oldval, key,))
while (self.l[i][0] is not oldval) or (self.l[i][1] is not key):
i += 1
self.l[i] = (val, key,)
else:
insort_left(self.l, (val, key,))
self.d[key] = val
assert self._assert_invariants()
return val
def remove(self, key, default=None, strictkey=True):
assert self._assert_invariants()
result = self.__delitem__(key, default, strictkey)
assert self._assert_invariants()
return result
def __getitem__(self, key, default=None, strictkey=True):
if not self.d.has_key(key):
if strictkey:
raise KeyError, key
else:
return default
return self.d[key]
def __delitem__(self, key, default=None, strictkey=True):
"""
@param strictkey: True if you want a KeyError in the case that
key is not there, False if you want a reference to default
in the case that key is not there
@param default: the object to return if key is not there; This
is ignored if strictkey.
@return: the object removed or default if there is not item by
that key and strictkey is False
"""
assert self._assert_invariants()
if self.d.has_key(key):
val = self.d.pop(key)
i = bisect_left(self.l, (val, key,))
while (self.l[i][0] is not val) or (self.l[i][1] is not key):
i += 1
self.l.pop(i)
assert self._assert_invariants()
return val
elif strictkey:
assert self._assert_invariants()
raise KeyError, key
else:
assert self._assert_invariants()
return default
def clear(self):
assert self._assert_invariants()
self.d.clear()
del self.l[:]
assert self._assert_invariants()
def update(self, otherdict):
"""
@return: self
"""
assert self._assert_invariants()
for (k, v,) in otherdict.iteritems():
self.insert(k, v)
assert self._assert_invariants()
return self
def has_key(self, key):
assert self._assert_invariants()
return self.d.has_key(key)
def popitem(self):
if not self.l:
raise KeyError, 'popitem(): dictionary is empty'
le = self.l.pop(0)
del self.d[le[1]]
return (le[1], le[0],)
def pop(self, k, default=None, strictkey=False):
if not self.d.has_key(k):
if strictkey:
raise KeyError, k
else:
return default
v = self.d.pop(k)
i = bisect_left(self.l, (v, k,))
while (self.l[i][0] is not v) or (self.l[i][1] is not k):
i += 1
self.l.pop(i)
return v
def pop_from_list(self, i=0):
le = self.l.pop(i)
del self.d[le[1]]
return le[1]
class AuxValueDict(dict):
"""I behave like a regular dict, but each key is associated with two
values: the main value, and an auxilliary one. Setting the main value
(with the usual d[key]=value) clears the auxvalue. You can set both main
and auxvalue at the same time, and can retrieve the values separately.
The main use case is a dictionary that represents unpacked child values
for a directory node, where a common pattern is to modify one or more
children and then pass the dict back to a packing function. The original
packed representation can be cached in the auxvalue, and the packing
function can use it directly on all unmodified children. On large
directories with a complex packing function, this can save considerable
time."""
def __init__(self, *args, **kwargs):
super(AuxValueDict, self).__init__(*args, **kwargs)
self.auxilliary = {}
def __setitem__(self, key, value):
super(AuxValueDict, self).__setitem__(key, value)
self.auxilliary[key] = None # clear the auxvalue
def __delitem__(self, key):
super(AuxValueDict, self).__delitem__(key)
self.auxilliary.pop(key)
def get_aux(self, key, default=None):
"""Retrieve the auxilliary value. There is no way to distinguish
between an auxvalue of 'None' and a key that does not have an
auxvalue, and get_aux() will not raise KeyError when called with a
missing key."""
return self.auxilliary.get(key, default)
def set_with_aux(self, key, value, auxilliary):
"""Set both the main value and the auxilliary value. There is no way
to distinguish between an auxvalue of 'None' and a key that does not
have an auxvalue."""
super(AuxValueDict, self).__setitem__(key, value)
self.auxilliary[key] = auxilliary