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