Lib/collections.py

   1 __all__ = ['Counter', 'deque', 'defaultdict', 'namedtuple', 'OrderedDict']
   2 # For bootstrapping reasons, the collection ABCs are defined in _abcoll.py.
   3 # They should however be considered an integral part of collections.py.
   4 from _abcoll import *
   5 import _abcoll
   6 __all__ += _abcoll.__all__
   7
   8 from _collections import deque, defaultdict
   9 from operator import itemgetter as _itemgetter, eq as _eq
  10 from keyword import iskeyword as _iskeyword
  11 import sys as _sys
  12 import heapq as _heapq
  13 from weakref import proxy as _proxy
  14 from itertools import repeat as _repeat, chain as _chain, starmap as _starmap, \
  15                       ifilter as _ifilter, imap as _imap, izip as _izip
  16
  17 ################################################################################
  18 ### OrderedDict
  19 ################################################################################
  20
  21 class _Link(object):
  22     __slots__ = 'prev', 'next', 'key', '__weakref__'
  23
  24 class OrderedDict(dict, MutableMapping):
  25     'Dictionary that remembers insertion order'
  26     # An inherited dict maps keys to values.
  27     # The inherited dict provides __getitem__, __len__, __contains__, and get.
  28     # The remaining methods are order-aware.
  29     # Big-O running times for all methods are the same as for regular dictionaries.
  30
  31     # The internal self.__map dictionary maps keys to links in a doubly linked list.
  32     # The circular doubly linked list starts and ends with a sentinel element.
  33     # The sentinel element never gets deleted (this simplifies the algorithm).
  34     # The prev/next links are weakref proxies (to prevent circular references).
  35     # Individual links are kept alive by the hard reference in self.__map.
  36     # Those hard references disappear when a key is deleted from an OrderedDict.
  37
  38     def __init__(self, *args, **kwds):
  39         if len(args) > 1:
  40             raise TypeError('expected at most 1 arguments, got %d' % len(args))
  41         try:
  42             self.__root
  43         except AttributeError:
  44             self.__root = root = _Link()    # sentinel node for the doubly linked list
  45             root.prev = root.next = root
  46             self.__map = {}
  47         self.update(*args, **kwds)
  48
  49     def clear(self):
  50         root = self.__root
  51         root.prev = root.next = root
  52         self.__map.clear()
  53         dict.clear(self)
  54
  55     def __setitem__(self, key, value):
  56         # Setting a new item creates a new link which goes at the end of the linked
  57         # list, and the inherited dictionary is updated with the new key/value pair.
  58         if key not in self:
  59             self.__map[key] = link = _Link()
  60             root = self.__root
  61             last = root.prev
  62             link.prev, link.next, link.key = last, root, key
  63             last.next = root.prev = _proxy(link)
  64         dict.__setitem__(self, key, value)
  65
  66     def __delitem__(self, key):
  67         # Deleting an existing item uses self.__map to find the link which is
  68         # then removed by updating the links in the predecessor and successor nodes.
  69         dict.__delitem__(self, key)
  70         link = self.__map.pop(key)
  71         link.prev.next = link.next
  72         link.next.prev = link.prev
  73
  74     def __iter__(self):
  75         # Traverse the linked list in order.
  76         root = self.__root
  77         curr = root.next
  78         while curr is not root:
  79             yield curr.key
  80             curr = curr.next
  81
  82     def __reversed__(self):
  83         # Traverse the linked list in reverse order.
  84         root = self.__root
  85         curr = root.prev
  86         while curr is not root:
  87             yield curr.key
  88             curr = curr.prev
  89
  90     def __reduce__(self):
  91         items = [[k, self[k]] for k in self]
  92         tmp = self.__map, self.__root
  93         del self.__map, self.__root
  94         inst_dict = vars(self).copy()
  95         self.__map, self.__root = tmp
  96         if inst_dict:
  97             return (self.__class__, (items,), inst_dict)
  98         return self.__class__, (items,)
  99
 100     setdefault = MutableMapping.setdefault
 101     update = MutableMapping.update
 102     pop = MutableMapping.pop
 103     keys = MutableMapping.keys
 104     values = MutableMapping.values
 105     items = MutableMapping.items
 106     iterkeys = MutableMapping.iterkeys
 107     itervalues = MutableMapping.itervalues
 108     iteritems = MutableMapping.iteritems
 109     __ne__ = MutableMapping.__ne__
 110
 111     def popitem(self, last=True):
 112         if not self:
 113             raise KeyError('dictionary is empty')
 114         key = next(reversed(self)) if last else next(iter(self))
 115         value = self.pop(key)
 116         return key, value
 117
 118     def __repr__(self):
 119         if not self:
 120             return '%s()' % (self.__class__.__name__,)
 121         return '%s(%r)' % (self.__class__.__name__, self.items())
 122
 123     def copy(self):
 124         return self.__class__(self)
 125
 126     @classmethod
 127     def fromkeys(cls, iterable, value=None):
 128         d = cls()
 129         for key in iterable:
 130             d[key] = value
 131         return d
 132
 133     def __eq__(self, other):
 134         if isinstance(other, OrderedDict):
 135             return len(self)==len(other) and \
 136                    all(_imap(_eq, self.iteritems(), other.iteritems()))
 137         return dict.__eq__(self, other)
 138
 139
 140
 141 ################################################################################
 142 ### namedtuple
 143 ################################################################################
 144
 145 def namedtuple(typename, field_names, verbose=False, rename=False):
 146     """Returns a new subclass of tuple with named fields.
 147
 148     >>> Point = namedtuple('Point', 'x y')
 149     >>> Point.__doc__                   # docstring for the new class
 150     'Point(x, y)'
 151     >>> p = Point(11, y=22)             # instantiate with positional args or keywords
 152     >>> p[0] + p[1]                     # indexable like a plain tuple
 153     33
 154     >>> x, y = p                        # unpack like a regular tuple
 155     >>> x, y
 156     (11, 22)
 157     >>> p.x + p.y                       # fields also accessable by name
 158     33
 159     >>> d = p._asdict()                 # convert to a dictionary
 160     >>> d['x']
 161     11
 162     >>> Point(**d)                      # convert from a dictionary
 163     Point(x=11, y=22)
 164     >>> p._replace(x=100)               # _replace() is like str.replace() but targets named fields
 165     Point(x=100, y=22)
 166
 167     """
 168
 169     # Parse and validate the field names.  Validation serves two purposes,
 170     # generating informative error messages and preventing template injection attacks.
 171     if isinstance(field_names, basestring):
 172         field_names = field_names.replace(',', ' ').split() # names separated by whitespace and/or commas
 173     field_names = tuple(map(str, field_names))
 174     if rename:
 175         names = list(field_names)
 176         seen = set()
 177         for i, name in enumerate(names):
 178             if (not all(c.isalnum() or c=='_' for c in name) or _iskeyword(name)
 179                 or not name or name[0].isdigit() or name.startswith('_')
 180                 or name in seen):
 181                 names[i] = '_%d' % (i+1)
 182             seen.add(name)
 183         field_names = tuple(names)
 184     for name in (typename,) + field_names:
 185         if not all(c.isalnum() or c=='_' for c in name):
 186             raise ValueError('Type names and field names can only contain alphanumeric characters and underscores: %r' % name)
 187         if _iskeyword(name):
 188             raise ValueError('Type names and field names cannot be a keyword: %r' % name)
 189         if name[0].isdigit():
 190             raise ValueError('Type names and field names cannot start with a number: %r' % name)
 191     seen_names = set()
 192     for name in field_names:
 193         if name.startswith('_') and not rename:
 194             raise ValueError('Field names cannot start with an underscore: %r' % name)
 195         if name in seen_names:
 196             raise ValueError('Encountered duplicate field name: %r' % name)
 197         seen_names.add(name)
 198
 199     # Create and fill-in the class template
 200     numfields = len(field_names)
 201     argtxt = repr(field_names).replace("'", "")[1:-1]   # tuple repr without parens or quotes
 202     reprtxt = ', '.join('%s=%%r' % name for name in field_names)
 203     template = '''class %(typename)s(tuple):
 204         '%(typename)s(%(argtxt)s)' \n
 205         __slots__ = () \n
 206         _fields = %(field_names)r \n
 207         def __new__(cls, %(argtxt)s):
 208             return tuple.__new__(cls, (%(argtxt)s)) \n
 209         @classmethod
 210         def _make(cls, iterable, new=tuple.__new__, len=len):
 211             'Make a new %(typename)s object from a sequence or iterable'
 212             result = new(cls, iterable)
 213             if len(result) != %(numfields)d:
 214                 raise TypeError('Expected %(numfields)d arguments, got %%d' %% len(result))
 215             return result \n
 216         def __repr__(self):
 217             return '%(typename)s(%(reprtxt)s)' %% self \n
 218         def _asdict(self):
 219             'Return a new OrderedDict which maps field names to their values'
 220             return OrderedDict(zip(self._fields, self)) \n
 221         def _replace(self, **kwds):
 222             'Return a new %(typename)s object replacing specified fields with new values'
 223             result = self._make(map(kwds.pop, %(field_names)r, self))
 224             if kwds:
 225                 raise ValueError('Got unexpected field names: %%r' %% kwds.keys())
 226             return result \n
 227         def __getnewargs__(self):
 228             return tuple(self) \n\n''' % locals()
 229     for i, name in enumerate(field_names):
 230         template += '        %s = property(itemgetter(%d))\n' % (name, i)
 231     if verbose:
 232         print template
 233
 234     # Execute the template string in a temporary namespace and
 235     # support tracing utilities by setting a value for frame.f_globals['__name__']
 236     namespace = dict(itemgetter=_itemgetter, __name__='namedtuple_%s' % typename,
 237                      OrderedDict=OrderedDict)
 238     try:
 239         exec template in namespace
 240     except SyntaxError, e:
 241         raise SyntaxError(e.message + ':\n' + template)
 242     result = namespace[typename]
 243
 244     # For pickling to work, the __module__ variable needs to be set to the frame
 245     # where the named tuple is created.  Bypass this step in enviroments where
 246     # sys._getframe is not defined (Jython for example).
 247     if hasattr(_sys, '_getframe'):
 248         result.__module__ = _sys._getframe(1).f_globals.get('__name__', '__main__')
 249
 250     return result
 251
 252
 253 ########################################################################
 254 ###  Counter
 255 ########################################################################
 256
 257 class Counter(dict):
 258     '''Dict subclass for counting hashable items.  Sometimes called a bag
 259     or multiset.  Elements are stored as dictionary keys and their counts
 260     are stored as dictionary values.
 261
 262     >>> c = Counter('abracadabra')      # count elements from a string
 263
 264     >>> c.most_common(3)                # three most common elements
 265     [('a', 5), ('r', 2), ('b', 2)]
 266     >>> sorted(c)                       # list all unique elements
 267     ['a', 'b', 'c', 'd', 'r']
 268     >>> ''.join(sorted(c.elements()))   # list elements with repetitions
 269     'aaaaabbcdrr'
 270     >>> sum(c.values())                 # total of all counts
 271     11
 272
 273     >>> c['a']                          # count of letter 'a'
 274     5
 275     >>> for elem in 'shazam':           # update counts from an iterable
 276     ...     c[elem] += 1                # by adding 1 to each element's count
 277     >>> c['a']                          # now there are seven 'a'
 278     7
 279     >>> del c['r']                      # remove all 'r'
 280     >>> c['r']                          # now there are zero 'r'
 281     0
 282
 283     >>> d = Counter('simsalabim')       # make another counter
 284     >>> c.update(d)                     # add in the second counter
 285     >>> c['a']                          # now there are nine 'a'
 286     9
 287
 288     >>> c.clear()                       # empty the counter
 289     >>> c
 290     Counter()
 291
 292     Note:  If a count is set to zero or reduced to zero, it will remain
 293     in the counter until the entry is deleted or the counter is cleared:
 294
 295     >>> c = Counter('aaabbc')
 296     >>> c['b'] -= 2                     # reduce the count of 'b' by two
 297     >>> c.most_common()                 # 'b' is still in, but its count is zero
 298     [('a', 3), ('c', 1), ('b', 0)]
 299
 300     '''
 301     # References:
 302     #   http://en.wikipedia.org/wiki/Multiset
 303     #   http://www.gnu.org/software/smalltalk/manual-base/html_node/Bag.html
 304     #   http://www.demo2s.com/Tutorial/Cpp/0380__set-multiset/Catalog0380__set-multiset.htm
 305     #   http://code.activestate.com/recipes/259174/
 306     #   Knuth, TAOCP Vol. II section 4.6.3
 307
 308     def __init__(self, iterable=None, **kwds):
 309         '''Create a new, empty Counter object.  And if given, count elements
 310         from an input iterable.  Or, initialize the count from another mapping
 311         of elements to their counts.
 312
 313         >>> c = Counter()                           # a new, empty counter
 314         >>> c = Counter('gallahad')                 # a new counter from an iterable
 315         >>> c = Counter({'a': 4, 'b': 2})           # a new counter from a mapping
 316         >>> c = Counter(a=4, b=2)                   # a new counter from keyword args
 317
 318         '''
 319         self.update(iterable, **kwds)
 320
 321     def __missing__(self, key):
 322         'The count of elements not in the Counter is zero.'
 323         # Needed so that self[missing_item] does not raise KeyError
 324         return 0
 325
 326     def most_common(self, n=None):
 327         '''List the n most common elements and their counts from the most
 328         common to the least.  If n is None, then list all element counts.
 329
 330         >>> Counter('abracadabra').most_common(3)
 331         [('a', 5), ('r', 2), ('b', 2)]
 332
 333         '''
 334         # Emulate Bag.sortedByCount from Smalltalk
 335         if n is None:
 336             return sorted(self.iteritems(), key=_itemgetter(1), reverse=True)
 337         return _heapq.nlargest(n, self.iteritems(), key=_itemgetter(1))
 338
 339     def elements(self):
 340         '''Iterator over elements repeating each as many times as its count.
 341
 342         >>> c = Counter('ABCABC')
 343         >>> sorted(c.elements())
 344         ['A', 'A', 'B', 'B', 'C', 'C']
 345
 346         # Knuth's example for prime factors of 1836:  2**2 * 3**3 * 17**1
 347         >>> prime_factors = Counter({2: 2, 3: 3, 17: 1})
 348         >>> product = 1
 349         >>> for factor in prime_factors.elements():     # loop over factors
 350         ...     product *= factor                       # and multiply them
 351         >>> product
 352         1836
 353
 354         Note, if an element's count has been set to zero or is a negative
 355         number, elements() will ignore it.
 356
 357         '''
 358         # Emulate Bag.do from Smalltalk and Multiset.begin from C++.
 359         return _chain.from_iterable(_starmap(_repeat, self.iteritems()))
 360
 361     # Override dict methods where necessary
 362
 363     @classmethod
 364     def fromkeys(cls, iterable, v=None):
 365         # There is no equivalent method for counters because setting v=1
 366         # means that no element can have a count greater than one.
 367         raise NotImplementedError(
 368             'Counter.fromkeys() is undefined.  Use Counter(iterable) instead.')
 369
 370     def update(self, iterable=None, **kwds):
 371         '''Like dict.update() but add counts instead of replacing them.
 372
 373         Source can be an iterable, a dictionary, or another Counter instance.
 374
 375         >>> c = Counter('which')
 376         >>> c.update('witch')           # add elements from another iterable
 377         >>> d = Counter('watch')
 378         >>> c.update(d)                 # add elements from another counter
 379         >>> c['h']                      # four 'h' in which, witch, and watch
 380         4
 381
 382         '''
 383         # The regular dict.update() operation makes no sense here because the
 384         # replace behavior results in the some of original untouched counts
 385         # being mixed-in with all of the other counts for a mismash that
 386         # doesn't have a straight-forward interpretation in most counting
 387         # contexts.  Instead, we implement straight-addition.  Both the inputs
 388         # and outputs are allowed to contain zero and negative counts.
 389
 390         if iterable is not None:
 391             if isinstance(iterable, Mapping):
 392                 if self:
 393                     for elem, count in iterable.iteritems():
 394                         self[elem] += count
 395                 else:
 396                     dict.update(self, iterable) # fast path when counter is empty
 397             else:
 398                 for elem in iterable:
 399                     self[elem] += 1
 400         if kwds:
 401             self.update(kwds)
 402
 403     def copy(self):
 404         'Like dict.copy() but returns a Counter instance instead of a dict.'
 405         return Counter(self)
 406
 407     def __delitem__(self, elem):
 408         'Like dict.__delitem__() but does not raise KeyError for missing values.'
 409         if elem in self:
 410             dict.__delitem__(self, elem)
 411
 412     def __repr__(self):
 413         if not self:
 414             return '%s()' % self.__class__.__name__
 415         items = ', '.join(map('%r: %r'.__mod__, self.most_common()))
 416         return '%s({%s})' % (self.__class__.__name__, items)
 417
 418     # Multiset-style mathematical operations discussed in:
 419     #       Knuth TAOCP Volume II section 4.6.3 exercise 19
 420     #       and at http://en.wikipedia.org/wiki/Multiset
 421     #
 422     # Outputs guaranteed to only include positive counts.
 423     #
 424     # To strip negative and zero counts, add-in an empty counter:
 425     #       c += Counter()
 426
 427     def __add__(self, other):
 428         '''Add counts from two counters.
 429
 430         >>> Counter('abbb') + Counter('bcc')
 431         Counter({'b': 4, 'c': 2, 'a': 1})
 432
 433         '''
 434         if not isinstance(other, Counter):
 435             return NotImplemented
 436         result = Counter()
 437         for elem in set(self) | set(other):
 438             newcount = self[elem] + other[elem]
 439             if newcount > 0:
 440                 result[elem] = newcount
 441         return result
 442
 443     def __sub__(self, other):
 444         ''' Subtract count, but keep only results with positive counts.
 445
 446         >>> Counter('abbbc') - Counter('bccd')
 447         Counter({'b': 2, 'a': 1})
 448
 449         '''
 450         if not isinstance(other, Counter):
 451             return NotImplemented
 452         result = Counter()
 453         for elem in set(self) | set(other):
 454             newcount = self[elem] - other[elem]
 455             if newcount > 0:
 456                 result[elem] = newcount
 457         return result
 458
 459     def __or__(self, other):
 460         '''Union is the maximum of value in either of the input counters.
 461
 462         >>> Counter('abbb') | Counter('bcc')
 463         Counter({'b': 3, 'c': 2, 'a': 1})
 464
 465         '''
 466         if not isinstance(other, Counter):
 467             return NotImplemented
 468         _max = max
 469         result = Counter()
 470         for elem in set(self) | set(other):
 471             newcount = _max(self[elem], other[elem])
 472             if newcount > 0:
 473                 result[elem] = newcount
 474         return result
 475
 476     def __and__(self, other):
 477         ''' Intersection is the minimum of corresponding counts.
 478
 479         >>> Counter('abbb') & Counter('bcc')
 480         Counter({'b': 1})
 481
 482         '''
 483         if not isinstance(other, Counter):
 484             return NotImplemented
 485         _min = min
 486         result = Counter()
 487         if len(self) < len(other):
 488             self, other = other, self
 489         for elem in _ifilter(self.__contains__, other):
 490             newcount = _min(self[elem], other[elem])
 491             if newcount > 0:
 492                 result[elem] = newcount
 493         return result
 494
 495
 496 if __name__ == '__main__':
 497     # verify that instances can be pickled
 498     from cPickle import loads, dumps
 499     Point = namedtuple('Point', 'x, y', True)
 500     p = Point(x=10, y=20)
 501     assert p == loads(dumps(p))
 502
 503     # test and demonstrate ability to override methods
 504     class Point(namedtuple('Point', 'x y')):
 505         __slots__ = ()
 506         @property
 507         def hypot(self):
 508             return (self.x ** 2 + self.y ** 2) ** 0.5
 509         def __str__(self):
 510             return 'Point: x=%6.3f  y=%6.3f  hypot=%6.3f' % (self.x, self.y, self.hypot)
 511
 512     for p in Point(3, 4), Point(14, 5/7.):
 513         print p
 514
 515     class Point(namedtuple('Point', 'x y')):
 516         'Point class with optimized _make() and _replace() without error-checking'
 517         __slots__ = ()
 518         _make = classmethod(tuple.__new__)
 519         def _replace(self, _map=map, **kwds):
 520             return self._make(_map(kwds.get, ('x', 'y'), self))
 521
 522     print Point(11, 22)._replace(x=100)
 523
 524     Point3D = namedtuple('Point3D', Point._fields + ('z',))
 525     print Point3D.__doc__
 526
 527     import doctest
 528     TestResults = namedtuple('TestResults', 'failed attempted')
 529     print TestResults(*doctest.testmod())