source file: /home/buildslave/tahoe/edgy/build/src/allmydata/util/dictutil.py
file stats: 416 lines, 413 executed: 99.3% covered
coverage versus previous test: 0 lines added, 0 lines removed
    1. """
    2. Tools to mess with dicts.
    3. """
    4. 
    5. import copy, operator
    6. from bisect import bisect_left, insort_left
    7. 
    8. from assertutil import _assert, precondition
    9. 
   10. def move(k, d1, d2, strict=False):
   11.     """
   12.     Move item with key k from d1 to d2.
   13.     """
   14.     if strict and not d1.has_key(k):
   15.         raise KeyError, k
   16. 
   17.     d2[k] = d1[k]
   18.     del d1[k]
   19. 
   20. def subtract(d1, d2):
   21.     """
   22.     Remove all items from d1 whose key occurs in d2.
   23. 
   24.     @returns d1
   25.     """
   26.     if len(d1) > len(d2):
   27.         for k in d2.keys():
   28.             if d1.has_key(k):
   29.                 del d1[k]
   30.     else:
   31.         for k in d1.keys():
   32.             if d2.has_key(k):
   33.                 del d1[k]
   34.     return d1
   35. 
   36. class DictOfSets(dict):
   37.     def add(self, key, value):
   38.         if key in self:
   39.             self[key].add(value)
   40.         else:
   41.             self[key] = set([value])
   42. 
   43.     def union(self, key, values):
   44.         if key in self:
   45.             self[key].update(values)
   46.         else:
   47.             self[key] = set(values)
   48. 
   49.     def update(self, otherdictofsets):
   50.         for key, value in otherdictofsets.iteritems():
   51.             self.union(key, value)
   52. 
   53.     def discard(self, key, value):
   54.         if not key in self:
   55.             return
   56.         self[key].discard(value)
   57.         if not self[key]:
   58.             del self[key]
   59. 
   60. class UtilDict:
   61.     def __init__(self, initialdata={}):
   62.         self.d = {}
   63.         self.update(initialdata)
   64. 
   65.     def del_if_present(self, key):
   66.         if self.has_key(key):
   67.             del self[key]
   68. 
   69.     def items_sorted_by_value(self):
   70.         """
   71.         @return a sequence of (key, value,) pairs sorted according to value
   72.         """
   73.         l = [(x[1], x[0],) for x in self.d.iteritems()]
   74.         l.sort()
   75.         return [(x[1], x[0],) for x in l]
   76. 
   77.     def items_sorted_by_key(self):
   78.         """
   79.         @return a sequence of (key, value,) pairs sorted according to key
   80.         """
   81.         l = self.d.items()
   82.         l.sort()
   83.         return l
   84. 
   85.     def __repr__(self, *args, **kwargs):
   86.         return self.d.__repr__(*args, **kwargs)
   87. 
   88.     def __str__(self, *args, **kwargs):
   89.         return self.d.__str__(*args, **kwargs)
   90. 
   91.     def __contains__(self, *args, **kwargs):
   92.         return self.d.__contains__(*args, **kwargs)
   93. 
   94.     def __len__(self, *args, **kwargs):
   95.         return self.d.__len__(*args, **kwargs)
   96. 
   97.     def __cmp__(self, other):
   98.         try:
   99.             return self.d.__cmp__(other)
  100.         except TypeError, le:
  101.             # 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")
  102.             try:
  103.                 return self.d.__cmp__(other.d)
  104.             except:
  105.                 raise le
  106. 
  107.     def __eq__(self, *args, **kwargs):
  108.         return self.d.__eq__(*args, **kwargs)
  109. 
  110.     def __ne__(self, *args, **kwargs):
  111.         return self.d.__ne__(*args, **kwargs)
  112. 
  113.     def __gt__(self, *args, **kwargs):
  114.         return self.d.__gt__(*args, **kwargs)
  115. 
  116.     def __ge__(self, *args, **kwargs):
  117.         return self.d.__ge__(*args, **kwargs)
  118. 
  119.     def __le__(self, *args, **kwargs):
  120.         return self.d.__le__(*args, **kwargs)
  121. 
  122.     def __lt__(self, *args, **kwargs):
  123.         return self.d.__lt__(*args, **kwargs)
  124. 
  125.     def __getitem__(self, *args, **kwargs):
  126.         return self.d.__getitem__(*args, **kwargs)
  127. 
  128.     def __setitem__(self, *args, **kwargs):
  129.         return self.d.__setitem__(*args, **kwargs)
  130. 
  131.     def __delitem__(self, *args, **kwargs):
  132.         return self.d.__delitem__(*args, **kwargs)
  133. 
  134.     def __iter__(self, *args, **kwargs):
  135.         return self.d.__iter__(*args, **kwargs)
  136. 
  137.     def clear(self, *args, **kwargs):
  138.         return self.d.clear(*args, **kwargs)
  139. 
  140.     def copy(self, *args, **kwargs):
  141.         return self.__class__(self.d.copy(*args, **kwargs))
  142. 
  143.     def fromkeys(self, *args, **kwargs):
  144.         return self.__class__(self.d.fromkeys(*args, **kwargs))
  145. 
  146.     def get(self, key, default=None):
  147.         return self.d.get(key, default)
  148. 
  149.     def has_key(self, *args, **kwargs):
  150.         return self.d.has_key(*args, **kwargs)
  151. 
  152.     def items(self, *args, **kwargs):
  153.         return self.d.items(*args, **kwargs)
  154. 
  155.     def iteritems(self, *args, **kwargs):
  156.         return self.d.iteritems(*args, **kwargs)
  157. 
  158.     def iterkeys(self, *args, **kwargs):
  159.         return self.d.iterkeys(*args, **kwargs)
  160. 
  161.     def itervalues(self, *args, **kwargs):
  162.         return self.d.itervalues(*args, **kwargs)
  163. 
  164.     def keys(self, *args, **kwargs):
  165.         return self.d.keys(*args, **kwargs)
  166. 
  167.     def pop(self, *args, **kwargs):
  168.         return self.d.pop(*args, **kwargs)
  169. 
  170.     def popitem(self, *args, **kwargs):
  171.         return self.d.popitem(*args, **kwargs)
  172. 
  173.     def setdefault(self, *args, **kwargs):
  174.         return self.d.setdefault(*args, **kwargs)
  175. 
  176.     def update(self, *args, **kwargs):
  177.         self.d.update(*args, **kwargs)
  178. 
  179.     def values(self, *args, **kwargs):
  180.         return self.d.values(*args, **kwargs)
  181. 
  182. class NumDict:
  183.     def __init__(self, initialdict={}):
  184.         self.d = copy.deepcopy(initialdict)
  185. 
  186.     def add_num(self, key, val, default=0):
  187.         """
  188.         If the key doesn't appear in self then it is created with value default
  189.         (before addition).
  190.         """
  191.         self.d[key] = self.d.get(key, default) + val
  192. 
  193.     def subtract_num(self, key, val, default=0):
  194.         self.d[key] = self.d.get(key, default) - val
  195. 
  196.     def sum(self):
  197.         """
  198.         @return: the sum of all values
  199.         """
  200.         return reduce(operator.__add__, self.d.values())
  201. 
  202.     def inc(self, key, default=0):
  203.         """
  204.         Increment the value associated with key in dict.  If there is no such
  205.         key, then one will be created with initial value 0 (before inc() --
  206.         therefore value 1 after inc).
  207.         """
  208.         self.add_num(key, 1, default)
  209. 
  210.     def dec(self, key, default=0):
  211.         """
  212.         Decrement the value associated with key in dict.  If there is no such
  213.         key, then one will be created with initial value 0 (before dec() --
  214.         therefore value -1 after dec).
  215.         """
  216.         self.subtract_num(key, 1, default)
  217. 
  218.     def items_sorted_by_value(self):
  219.         """
  220.         @return a sequence of (key, value,) pairs sorted according to value
  221.         """
  222.         l = [(x[1], x[0],) for x in self.d.iteritems()]
  223.         l.sort()
  224.         return [(x[1], x[0],) for x in l]
  225. 
  226.     def item_with_largest_value(self):
  227.         it = self.d.iteritems()
  228.         (winner, winnerval,) = it.next()
  229.         try:
  230.             while True:
  231.                 n, nv = it.next()
  232.                 if nv > winnerval:
  233.                     winner = n
  234.                     winnerval = nv
  235.         except StopIteration:
  236.             pass
  237.         return (winner, winnerval,)
  238. 
  239.     def items_sorted_by_key(self):
  240.         """
  241.         @return a sequence of (key, value,) pairs sorted according to key
  242.         """
  243.         l = self.d.items()
  244.         l.sort()
  245.         return l
  246. 
  247.     def __repr__(self, *args, **kwargs):
  248.         return self.d.__repr__(*args, **kwargs)
  249. 
  250.     def __str__(self, *args, **kwargs):
  251.         return self.d.__str__(*args, **kwargs)
  252. 
  253.     def __contains__(self, *args, **kwargs):
  254.         return self.d.__contains__(*args, **kwargs)
  255. 
  256.     def __len__(self, *args, **kwargs):
  257.         return self.d.__len__(*args, **kwargs)
  258. 
  259.     def __cmp__(self, other):
  260.         try:
  261.             return self.d.__cmp__(other)
  262.         except TypeError, le:
  263.             # 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")
  264.             try:
  265.                 return self.d.__cmp__(other.d)
  266.             except:
  267.                 raise le
  268. 
  269.     def __eq__(self, *args, **kwargs):
  270.         return self.d.__eq__(*args, **kwargs)
  271. 
  272.     def __ne__(self, *args, **kwargs):
  273.         return self.d.__ne__(*args, **kwargs)
  274. 
  275.     def __gt__(self, *args, **kwargs):
  276.         return self.d.__gt__(*args, **kwargs)
  277. 
  278.     def __ge__(self, *args, **kwargs):
  279.         return self.d.__ge__(*args, **kwargs)
  280. 
  281.     def __le__(self, *args, **kwargs):
  282.         return self.d.__le__(*args, **kwargs)
  283. 
  284.     def __lt__(self, *args, **kwargs):
  285.         return self.d.__lt__(*args, **kwargs)
  286. 
  287.     def __getitem__(self, *args, **kwargs):
  288.         return self.d.__getitem__(*args, **kwargs)
  289. 
  290.     def __setitem__(self, *args, **kwargs):
  291.         return self.d.__setitem__(*args, **kwargs)
  292. 
  293.     def __delitem__(self, *args, **kwargs):
  294.         return self.d.__delitem__(*args, **kwargs)
  295. 
  296.     def __iter__(self, *args, **kwargs):
  297.         return self.d.__iter__(*args, **kwargs)
  298. 
  299.     def clear(self, *args, **kwargs):
  300.         return self.d.clear(*args, **kwargs)
  301. 
  302.     def copy(self, *args, **kwargs):
  303.         return self.__class__(self.d.copy(*args, **kwargs))
  304. 
  305.     def fromkeys(self, *args, **kwargs):
  306.         return self.__class__(self.d.fromkeys(*args, **kwargs))
  307. 
  308.     def get(self, key, default=0):
  309.         return self.d.get(key, default)
  310. 
  311.     def has_key(self, *args, **kwargs):
  312.         return self.d.has_key(*args, **kwargs)
  313. 
  314.     def items(self, *args, **kwargs):
  315.         return self.d.items(*args, **kwargs)
  316. 
  317.     def iteritems(self, *args, **kwargs):
  318.         return self.d.iteritems(*args, **kwargs)
  319. 
  320.     def iterkeys(self, *args, **kwargs):
  321.         return self.d.iterkeys(*args, **kwargs)
  322. 
  323.     def itervalues(self, *args, **kwargs):
  324.         return self.d.itervalues(*args, **kwargs)
  325. 
  326.     def keys(self, *args, **kwargs):
  327.         return self.d.keys(*args, **kwargs)
  328. 
  329.     def pop(self, *args, **kwargs):
  330.         return self.d.pop(*args, **kwargs)
  331. 
  332.     def popitem(self, *args, **kwargs):
  333.         return self.d.popitem(*args, **kwargs)
  334. 
  335.     def setdefault(self, *args, **kwargs):
  336.         return self.d.setdefault(*args, **kwargs)
  337. 
  338.     def update(self, *args, **kwargs):
  339.         return self.d.update(*args, **kwargs)
  340. 
  341.     def values(self, *args, **kwargs):
  342.         return self.d.values(*args, **kwargs)
  343. 
  344. def del_if_present(d, k):
  345.     if d.has_key(k):
  346.         del d[k]
  347. 
  348. class ValueOrderedDict:
  349.     """
  350.     Note: this implementation assumes that the values do not mutate and change
  351.     their sort order.  That is, it stores the values in a sorted list and
  352.     as items are added and removed from the dict, it makes updates to the list
  353.     which will keep the list sorted.  But if a value that is currently sitting
  354.     in the list changes its sort order, then the internal consistency of this
  355.     object will be lost.
  356. 
  357.     If that happens, and if assertion checking is turned on, then you will get
  358.     an assertion failure the very next time you try to do anything with this
  359.     ValueOrderedDict.  However, those internal consistency checks are very slow
  360.     and almost certainly unacceptable to leave turned on in production code.
  361.     """
  362.     class ItemIterator:
  363.         def __init__(self, c):
  364.             self.c = c
  365.             self.i = 0
  366.         def __iter__(self):
  367.             return self
  368.         def next(self):
  369.             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)
  370.             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)
  371.             if self.i == len(self.c.l):
  372.                 raise StopIteration
  373.             le = self.c.l[self.i]
  374.             self.i += 1
  375.             return (le[1], le[0],)
  376. 
  377.     def iteritems(self):
  378.         return ValueOrderedDict.ItemIterator(self)
  379. 
  380.     def items(self):
  381.         return zip(map(operator.__getitem__, self.l, [1]*len(self.l)), map(operator.__getitem__, self.l, [0]*len(self.l)))
  382. 
  383.     def values(self):
  384.         return map(operator.__getitem__, self.l, [0]*len(self.l))
  385. 
  386.     def keys(self):
  387.         return map(operator.__getitem__, self.l, [1]*len(self.l))
  388. 
  389.     class KeyIterator:
  390.         def __init__(self, c):
  391.             self.c = c
  392.             self.i = 0
  393.         def __iter__(self):
  394.             return self
  395.         def next(self):
  396.             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)
  397.             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)
  398.             if self.i == len(self.c.l):
  399.                 raise StopIteration
  400.             le = self.c.l[self.i]
  401.             self.i += 1
  402.             return le[1]
  403. 
  404.     def iterkeys(self):
  405.         return ValueOrderedDict.KeyIterator(self)
  406. 
  407.     class ValueIterator:
  408.         def __init__(self, c):
  409.             self.c = c
  410.             self.i = 0
  411.         def __iter__(self):
  412.             return self
  413.         def next(self):
  414.             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)
  415.             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)
  416.             if self.i == len(self.c.l):
  417.                 raise StopIteration
  418.             le = self.c.l[self.i]
  419.             self.i += 1
  420.             return le[0]
  421. 
  422.     def itervalues(self):
  423.         return ValueOrderedDict.ValueIterator(self)
  424. 
  425.     def __init__(self, initialdata={}):
  426.         self.d = {} # k: key, v: val
  427.         self.l = [] # sorted list of tuples of (val, key,)
  428.         self.update(initialdata)
  429.         assert self._assert_invariants()
  430. 
  431.     def __len__(self):
  432.         return len(self.l)
  433. 
  434.     def __repr_n__(self, n=None):
  435.         s = ["{",]
  436.         try:
  437.             iter = self.iteritems()
  438.             x = iter.next()
  439.             s.append(str(x[0])); s.append(": "); s.append(str(x[1]))
  440.             i = 1
  441.             while (n is None) or (i < n):
  442.                 x = iter.next()
  443.                 s.append(", "); s.append(str(x[0])); s.append(": "); s.append(str(x[1]))
  444.         except StopIteration:
  445.             pass
  446.         s.append("}")
  447.         return ''.join(s)
  448. 
  449.     def __repr__(self):
  450.         return "<%s %s>" % (self.__class__.__name__, self.__repr_n__(),)
  451. 
  452.     def __str__(self):
  453.         return "<%s %s>" % (self.__class__.__name__, self.__repr_n__(16),)
  454. 
  455.     def __eq__(self, other):
  456.         for (k, v,) in other.iteritems():
  457.             if not self.d.has_key(k) or self.d[k] != v:
  458.                 return False
  459.         return True
  460. 
  461.     def __ne__(self, other):
  462.         return not self.__eq__(other)
  463. 
  464.     def _assert_invariants(self):
  465.         iter = self.l.__iter__()
  466.         try:
  467.             oldx = iter.next()
  468.             while True:
  469.                 x = iter.next()
  470.                 # self.l is required to be sorted
  471.                 _assert(x >= oldx, x, oldx)
  472.                 # every element of self.l is required to appear in self.d
  473.                 _assert(self.d.has_key(x[1]), x)
  474.                 oldx =x
  475.         except StopIteration:
  476.             pass
  477.         for (k, v,) in self.d.iteritems():
  478.             i = bisect_left(self.l, (v, k,))
  479.             while (self.l[i][0] is not v) or (self.l[i][1] is not k):
  480.                 i += 1
  481.             _assert(i < len(self.l), i, len(self.l), k, v, self.l)
  482.             _assert(self.l[i][0] is v, i, v, l=self.l, d=self.d)
  483.             _assert(self.l[i][1] is k, i, k, l=self.l, d=self.d)
  484.         return True
  485. 
  486.     def insert(self, key, val=None):
  487.         assert self._assert_invariants()
  488.         result = self.__setitem__(key, val)
  489.         assert self._assert_invariants()
  490.         return result
  491. 
  492.     def setdefault(self, key, default=None):
  493.         assert self._assert_invariants()
  494.         if not self.has_key(key):
  495.             self[key] = default
  496.         assert self._assert_invariants()
  497.         return self[key]
  498. 
  499.     def __setitem__(self, key, val=None):
  500.         assert self._assert_invariants()
  501.         if self.d.has_key(key):
  502.             oldval = self.d[key]
  503.             if oldval != val:
  504.                 # re-sort
  505.                 i = bisect_left(self.l, (oldval, key,))
  506.                 while (self.l[i][0] is not oldval) or (self.l[i][1] is not key):
  507.                     i += 1
  508.                 self.l.pop(i)
  509.                 insort_left(self.l, (val, key,))
  510.             elif oldval is not val:
  511.                 # replace
  512.                 i = bisect_left(self.l, (oldval, key,))
  513.                 while (self.l[i][0] is not oldval) or (self.l[i][1] is not key):
  514.                     i += 1
  515.                 self.l[i] = (val, key,)
  516.         else:
  517.             insort_left(self.l, (val, key,))
  518. 
  519.         self.d[key] = val
  520.         assert self._assert_invariants()
  521.         return val
  522. 
  523.     def remove(self, key, default=None, strictkey=True):
  524.         assert self._assert_invariants()
  525.         result = self.__delitem__(key, default, strictkey)
  526.         assert self._assert_invariants()
  527.         return result
  528. 
  529.     def __getitem__(self, key, default=None, strictkey=True):
  530.         if not self.d.has_key(key):
  531.             if strictkey:
  532.                 raise KeyError, key
  533.             else:
  534.                 return default
  535.         return self.d[key]
  536. 
  537.     def __delitem__(self, key, default=None, strictkey=True):
  538.         """
  539.         @param strictkey: True if you want a KeyError in the case that
  540.             key is not there, False if you want a reference to default
  541.             in the case that key is not there
  542.         @param default: the object to return if key is not there; This
  543.             is ignored if strictkey.
  544. 
  545.         @return: the object removed or default if there is not item by
  546.             that key and strictkey is False
  547.         """
  548.         assert self._assert_invariants()
  549.         if self.d.has_key(key):
  550.             val = self.d.pop(key)
  551.             i = bisect_left(self.l, (val, key,))
  552.             while (self.l[i][0] is not val) or (self.l[i][1] is not key):
  553.                 i += 1
  554.             self.l.pop(i)
  555.             assert self._assert_invariants()
  556.             return val
  557.         elif strictkey:
  558.             assert self._assert_invariants()
  559.             raise KeyError, key
  560.         else:
  561.             assert self._assert_invariants()
  562.             return default
  563. 
  564.     def clear(self):
  565.         assert self._assert_invariants()
  566.         self.d.clear()
  567.         del self.l[:]
  568.         assert self._assert_invariants()
  569. 
  570.     def update(self, otherdict):
  571.         """
  572.         @return: self
  573.         """
  574.         assert self._assert_invariants()
  575.         for (k, v,) in otherdict.iteritems():
  576.             self.insert(k, v)
  577.         assert self._assert_invariants()
  578.         return self
  579. 
  580.     def has_key(self, key):
  581.         assert self._assert_invariants()
  582.         return self.d.has_key(key)
  583. 
  584.     def popitem(self):
  585.         if not self.l:
  586.             raise KeyError, 'popitem(): dictionary is empty'
  587.         le = self.l.pop(0)
  588.         del self.d[le[1]]
  589.         return (le[1], le[0],)
  590. 
  591.     def pop(self, k, default=None, strictkey=False):
  592.         if not self.d.has_key(k):
  593.             if strictkey:
  594.                 raise KeyError, k
  595.             else:
  596.                 return default
  597.         v = self.d.pop(k)
  598.         i = bisect_left(self.l, (v, k,))
  599.         while (self.l[i][0] is not v) or (self.l[i][1] is not k):
  600.             i += 1
  601.         self.l.pop(i)
  602.         return v
  603. 
  604.     def pop_from_list(self, i=0):
  605.         le = self.l.pop(i)
  606.         del self.d[le[1]]
  607.         return le[1]