src/allmydata/util/dictutil.py

   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]
 608
 609 class AuxValueDict(dict):
 610     """I behave like a regular dict, but each key is associated with two
 611     values: the main value, and an auxilliary one. Setting the main value
 612     (with the usual d[key]=value) clears the auxvalue. You can set both main
 613     and auxvalue at the same time, and can retrieve the values separately.
 614
 615     The main use case is a dictionary that represents unpacked child values
 616     for a directory node, where a common pattern is to modify one or more
 617     children and then pass the dict back to a packing function. The original
 618     packed representation can be cached in the auxvalue, and the packing
 619     function can use it directly on all unmodified children. On large
 620     directories with a complex packing function, this can save considerable
 621     time."""
 622
 623     def __init__(self, *args, **kwargs):
 624         super(AuxValueDict, self).__init__(*args, **kwargs)
 625         self.auxilliary = {}
 626
 627     def __setitem__(self, key, value):
 628         super(AuxValueDict, self).__setitem__(key, value)
 629         self.auxilliary[key] = None # clear the auxvalue
 630
 631     def __delitem__(self, key):
 632         super(AuxValueDict, self).__delitem__(key)
 633         self.auxilliary.pop(key)
 634
 635     def get_aux(self, key, default=None):
 636         """Retrieve the auxilliary value. There is no way to distinguish
 637         between an auxvalue of 'None' and a key that does not have an
 638         auxvalue, and get_aux() will not raise KeyError when called with a
 639         missing key."""
 640         return self.auxilliary.get(key, default)
 641
 642     def set_with_aux(self, key, value, auxilliary):
 643         """Set both the main value and the auxilliary value. There is no way
 644         to distinguish between an auxvalue of 'None' and a key that does not
 645         have an auxvalue."""
 646         super(AuxValueDict, self).__setitem__(key, value)
 647         self.auxilliary[key] = auxilliary