blob: 8fb123fbc107be81a9d03396d846f16e95ea92ea [file] [log] [blame]
Brett Cannon23a4a7b2008-05-12 00:56:28 +00001__all__ = ['deque', 'defaultdict', 'namedtuple', 'UserDict', 'UserList',
Raymond Hettinger2d32f632009-03-02 21:24:57 +00002 'UserString', 'Counter', 'OrderedDict']
Guido van Rossumcd16bf62007-06-13 18:07:49 +00003# For bootstrapping reasons, the collection ABCs are defined in _abcoll.py.
4# They should however be considered an integral part of collections.py.
5from _abcoll import *
6import _abcoll
7__all__ += _abcoll.__all__
8
Christian Heimes99170a52007-12-19 02:07:34 +00009from _collections import deque, defaultdict
10from operator import itemgetter as _itemgetter
11from keyword import iskeyword as _iskeyword
12import sys as _sys
Raymond Hettingerb8baf632009-01-14 02:20:07 +000013import heapq as _heapq
Raymond Hettinger798ee1a2009-03-23 18:29:11 +000014from weakref import proxy as _proxy
Raymond Hettingerea9f8db2009-03-02 21:28:41 +000015from itertools import repeat as _repeat, chain as _chain, starmap as _starmap
Raymond Hettinger2d32f632009-03-02 21:24:57 +000016
17################################################################################
18### OrderedDict
19################################################################################
20
Raymond Hettingerf1736542009-03-23 05:19:21 +000021class _Link(object):
22 __slots__ = 'prev', 'next', 'key', '__weakref__'
23
Raymond Hettinger2d32f632009-03-02 21:24:57 +000024class OrderedDict(dict, MutableMapping):
Raymond Hettinger18ed2cb2009-03-19 23:14:39 +000025 'Dictionary that remembers insertion order'
Raymond Hettingerf1736542009-03-23 05:19:21 +000026 # An inherited dict maps keys to values.
Raymond Hettinger18ed2cb2009-03-19 23:14:39 +000027 # The inherited dict provides __getitem__, __len__, __contains__, and get.
28 # The remaining methods are order-aware.
Raymond Hettingerf1736542009-03-23 05:19:21 +000029 # 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.
Raymond Hettinger2d32f632009-03-02 21:24:57 +000037
38 def __init__(self, *args, **kwds):
39 if len(args) > 1:
40 raise TypeError('expected at most 1 arguments, got %d' % len(args))
Raymond Hettinger08c70cf2009-03-03 20:47:29 +000041 try:
Raymond Hettingerf1736542009-03-23 05:19:21 +000042 self.__root
Raymond Hettinger08c70cf2009-03-03 20:47:29 +000043 except AttributeError:
Raymond Hettinger52dc06b2009-03-25 22:45:22 +000044 self.__root = root = _Link() # sentinel node for the doubly linked list
45 root.prev = root.next = root
46 self.__map = {}
Raymond Hettinger2d32f632009-03-02 21:24:57 +000047 self.update(*args, **kwds)
48
49 def clear(self):
Raymond Hettingerf1736542009-03-23 05:19:21 +000050 root = self.__root
51 root.prev = root.next = root
Raymond Hettinger52dc06b2009-03-25 22:45:22 +000052 self.__map.clear()
Raymond Hettinger2d32f632009-03-02 21:24:57 +000053 dict.clear(self)
54
55 def __setitem__(self, key, value):
Raymond Hettingerf1736542009-03-23 05:19:21 +000056 # 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.
Raymond Hettinger2d32f632009-03-02 21:24:57 +000058 if key not in self:
Raymond Hettingerf1736542009-03-23 05:19:21 +000059 self.__map[key] = link = _Link()
60 root = self.__root
61 last = root.prev
62 link.prev, link.next, link.key = last, root, key
Raymond Hettinger798ee1a2009-03-23 18:29:11 +000063 last.next = root.prev = _proxy(link)
Raymond Hettinger2d32f632009-03-02 21:24:57 +000064 dict.__setitem__(self, key, value)
65
66 def __delitem__(self, key):
Raymond Hettingerf1736542009-03-23 05:19:21 +000067 # 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.
Raymond Hettinger2d32f632009-03-02 21:24:57 +000069 dict.__delitem__(self, key)
Raymond Hettingerf1736542009-03-23 05:19:21 +000070 link = self.__map.pop(key)
71 link.prev.next = link.next
72 link.next.prev = link.prev
Raymond Hettinger2d32f632009-03-02 21:24:57 +000073
74 def __iter__(self):
Raymond Hettingerf1736542009-03-23 05:19:21 +000075 # 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
Raymond Hettinger2d32f632009-03-02 21:24:57 +000081
82 def __reversed__(self):
Raymond Hettingerf1736542009-03-23 05:19:21 +000083 # 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
Raymond Hettinger2d32f632009-03-02 21:24:57 +000089
90 def __reduce__(self):
91 items = [[k, self[k]] for k in self]
Raymond Hettingerf1736542009-03-23 05:19:21 +000092 tmp = self.__map, self.__root
93 del self.__map, self.__root
Raymond Hettinger2d32f632009-03-02 21:24:57 +000094 inst_dict = vars(self).copy()
Raymond Hettingerf1736542009-03-23 05:19:21 +000095 self.__map, self.__root = tmp
Raymond Hettinger14b89ff2009-03-03 22:20:56 +000096 if inst_dict:
97 return (self.__class__, (items,), inst_dict)
98 return self.__class__, (items,)
Raymond Hettinger2d32f632009-03-02 21:24:57 +000099
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
Raymond Hettinger18ed2cb2009-03-19 23:14:39 +0000107 def popitem(self, last=True):
108 if not self:
109 raise KeyError('dictionary is empty')
110 key = next(reversed(self)) if last else next(iter(self))
111 value = self.pop(key)
112 return key, value
113
Raymond Hettinger2d32f632009-03-02 21:24:57 +0000114 def __repr__(self):
115 if not self:
116 return '%s()' % (self.__class__.__name__,)
117 return '%s(%r)' % (self.__class__.__name__, list(self.items()))
118
119 def copy(self):
120 return self.__class__(self)
121
122 @classmethod
123 def fromkeys(cls, iterable, value=None):
124 d = cls()
125 for key in iterable:
126 d[key] = value
127 return d
128
129 def __eq__(self, other):
130 if isinstance(other, OrderedDict):
Raymond Hettinger798ee1a2009-03-23 18:29:11 +0000131 return len(self)==len(other) and \
132 all(p==q for p, q in zip(self.items(), other.items()))
Raymond Hettinger2d32f632009-03-02 21:24:57 +0000133 return dict.__eq__(self, other)
134
Benjamin Peterson2504b7a2009-04-04 17:26:32 +0000135 def __ne__(self, other):
136 return not self == other
137
Raymond Hettingerb8baf632009-01-14 02:20:07 +0000138
Christian Heimes99170a52007-12-19 02:07:34 +0000139
Raymond Hettinger48b8b662008-02-05 01:53:00 +0000140################################################################################
141### namedtuple
142################################################################################
143
Benjamin Petersona86f2c02009-02-10 02:41:10 +0000144def namedtuple(typename, field_names, verbose=False, rename=False):
Guido van Rossumd8faa362007-04-27 19:54:29 +0000145 """Returns a new subclass of tuple with named fields.
146
Guido van Rossum8ce8a782007-11-01 19:42:39 +0000147 >>> Point = namedtuple('Point', 'x y')
Thomas Wouters1b7f8912007-09-19 03:06:30 +0000148 >>> Point.__doc__ # docstring for the new class
Guido van Rossumd8faa362007-04-27 19:54:29 +0000149 'Point(x, y)'
Thomas Wouters1b7f8912007-09-19 03:06:30 +0000150 >>> p = Point(11, y=22) # instantiate with positional args or keywords
Christian Heimes99170a52007-12-19 02:07:34 +0000151 >>> p[0] + p[1] # indexable like a plain tuple
Guido van Rossumd8faa362007-04-27 19:54:29 +0000152 33
Christian Heimes99170a52007-12-19 02:07:34 +0000153 >>> x, y = p # unpack like a regular tuple
Guido van Rossumd8faa362007-04-27 19:54:29 +0000154 >>> x, y
155 (11, 22)
Thomas Wouters1b7f8912007-09-19 03:06:30 +0000156 >>> p.x + p.y # fields also accessable by name
Guido van Rossumd8faa362007-04-27 19:54:29 +0000157 33
Christian Heimes0449f632007-12-15 01:27:15 +0000158 >>> d = p._asdict() # convert to a dictionary
Guido van Rossum8ce8a782007-11-01 19:42:39 +0000159 >>> d['x']
160 11
161 >>> Point(**d) # convert from a dictionary
Guido van Rossumd8faa362007-04-27 19:54:29 +0000162 Point(x=11, y=22)
Christian Heimes0449f632007-12-15 01:27:15 +0000163 >>> p._replace(x=100) # _replace() is like str.replace() but targets named fields
Thomas Wouters1b7f8912007-09-19 03:06:30 +0000164 Point(x=100, y=22)
Guido van Rossumd8faa362007-04-27 19:54:29 +0000165
166 """
167
Christian Heimes2380ac72008-01-09 00:17:24 +0000168 # Parse and validate the field names. Validation serves two purposes,
169 # generating informative error messages and preventing template injection attacks.
Guido van Rossum8ce8a782007-11-01 19:42:39 +0000170 if isinstance(field_names, str):
171 field_names = field_names.replace(',', ' ').split() # names separated by whitespace and/or commas
Benjamin Petersone9bbc8b2008-09-28 02:06:32 +0000172 field_names = tuple(map(str, field_names))
Benjamin Petersona86f2c02009-02-10 02:41:10 +0000173 if rename:
174 names = list(field_names)
175 seen = set()
176 for i, name in enumerate(names):
177 if (not all(c.isalnum() or c=='_' for c in name) or _iskeyword(name)
178 or not name or name[0].isdigit() or name.startswith('_')
179 or name in seen):
Raymond Hettinger56145242009-04-02 22:31:59 +0000180 names[i] = '_%d' % i
Benjamin Petersona86f2c02009-02-10 02:41:10 +0000181 seen.add(name)
182 field_names = tuple(names)
Guido van Rossum8ce8a782007-11-01 19:42:39 +0000183 for name in (typename,) + field_names:
Christian Heimesb9eccbf2007-12-05 20:18:38 +0000184 if not all(c.isalnum() or c=='_' for c in name):
Guido van Rossum8ce8a782007-11-01 19:42:39 +0000185 raise ValueError('Type names and field names can only contain alphanumeric characters and underscores: %r' % name)
186 if _iskeyword(name):
187 raise ValueError('Type names and field names cannot be a keyword: %r' % name)
188 if name[0].isdigit():
189 raise ValueError('Type names and field names cannot start with a number: %r' % name)
190 seen_names = set()
191 for name in field_names:
Benjamin Petersona86f2c02009-02-10 02:41:10 +0000192 if name.startswith('_') and not rename:
Christian Heimes0449f632007-12-15 01:27:15 +0000193 raise ValueError('Field names cannot start with an underscore: %r' % name)
Guido van Rossum8ce8a782007-11-01 19:42:39 +0000194 if name in seen_names:
195 raise ValueError('Encountered duplicate field name: %r' % name)
196 seen_names.add(name)
197
198 # Create and fill-in the class template
Christian Heimesfaf2f632008-01-06 16:59:19 +0000199 numfields = len(field_names)
Thomas Wouters1b7f8912007-09-19 03:06:30 +0000200 argtxt = repr(field_names).replace("'", "")[1:-1] # tuple repr without parens or quotes
Guido van Rossumd59da4b2007-05-22 18:11:13 +0000201 reprtxt = ', '.join('%s=%%r' % name for name in field_names)
202 template = '''class %(typename)s(tuple):
Christian Heimes0449f632007-12-15 01:27:15 +0000203 '%(typename)s(%(argtxt)s)' \n
204 __slots__ = () \n
Christian Heimesfaf2f632008-01-06 16:59:19 +0000205 _fields = %(field_names)r \n
Guido van Rossumd59da4b2007-05-22 18:11:13 +0000206 def __new__(cls, %(argtxt)s):
Christian Heimes0449f632007-12-15 01:27:15 +0000207 return tuple.__new__(cls, (%(argtxt)s)) \n
Christian Heimesfaf2f632008-01-06 16:59:19 +0000208 @classmethod
Christian Heimes043d6f62008-01-07 17:19:16 +0000209 def _make(cls, iterable, new=tuple.__new__, len=len):
Christian Heimesfaf2f632008-01-06 16:59:19 +0000210 'Make a new %(typename)s object from a sequence or iterable'
Christian Heimes043d6f62008-01-07 17:19:16 +0000211 result = new(cls, iterable)
Christian Heimesfaf2f632008-01-06 16:59:19 +0000212 if len(result) != %(numfields)d:
213 raise TypeError('Expected %(numfields)d arguments, got %%d' %% len(result))
214 return result \n
Guido van Rossumd59da4b2007-05-22 18:11:13 +0000215 def __repr__(self):
Christian Heimes0449f632007-12-15 01:27:15 +0000216 return '%(typename)s(%(reprtxt)s)' %% self \n
Raymond Hettingera4f52b12009-03-02 22:28:31 +0000217 def _asdict(self):
218 'Return a new OrderedDict which maps field names to their values'
219 return OrderedDict(zip(self._fields, self)) \n
Christian Heimes0449f632007-12-15 01:27:15 +0000220 def _replace(self, **kwds):
Guido van Rossum3d392eb2007-11-16 00:35:22 +0000221 'Return a new %(typename)s object replacing specified fields with new values'
Christian Heimesfaf2f632008-01-06 16:59:19 +0000222 result = self._make(map(kwds.pop, %(field_names)r, self))
223 if kwds:
224 raise ValueError('Got unexpected field names: %%r' %% kwds.keys())
Georg Brandlc28e1fa2008-06-10 19:20:26 +0000225 return result \n
226 def __getnewargs__(self):
227 return tuple(self) \n\n''' % locals()
Guido van Rossumd59da4b2007-05-22 18:11:13 +0000228 for i, name in enumerate(field_names):
Thomas Wouters1b7f8912007-09-19 03:06:30 +0000229 template += ' %s = property(itemgetter(%d))\n' % (name, i)
230 if verbose:
231 print(template)
Guido van Rossum8ce8a782007-11-01 19:42:39 +0000232
Georg Brandlf08a9dd2008-06-10 16:57:31 +0000233 # Execute the template string in a temporary namespace and
234 # support tracing utilities by setting a value for frame.f_globals['__name__']
Raymond Hettingera4f52b12009-03-02 22:28:31 +0000235 namespace = dict(itemgetter=_itemgetter, __name__='namedtuple_%s' % typename,
236 OrderedDict=OrderedDict)
Guido van Rossum8ce8a782007-11-01 19:42:39 +0000237 try:
238 exec(template, namespace)
239 except SyntaxError as e:
Christian Heimes99170a52007-12-19 02:07:34 +0000240 raise SyntaxError(e.msg + ':\n' + template) from e
Guido van Rossum8ce8a782007-11-01 19:42:39 +0000241 result = namespace[typename]
242
243 # For pickling to work, the __module__ variable needs to be set to the frame
244 # where the named tuple is created. Bypass this step in enviroments where
245 # sys._getframe is not defined (Jython for example).
Guido van Rossumd59da4b2007-05-22 18:11:13 +0000246 if hasattr(_sys, '_getframe'):
Raymond Hettinger0f055172009-01-27 10:06:09 +0000247 result.__module__ = _sys._getframe(1).f_globals.get('__name__', '__main__')
Guido van Rossum8ce8a782007-11-01 19:42:39 +0000248
Guido van Rossumd59da4b2007-05-22 18:11:13 +0000249 return result
Guido van Rossumd8faa362007-04-27 19:54:29 +0000250
Guido van Rossumd8faa362007-04-27 19:54:29 +0000251
Raymond Hettingerb8baf632009-01-14 02:20:07 +0000252########################################################################
253### Counter
254########################################################################
255
256class Counter(dict):
257 '''Dict subclass for counting hashable items. Sometimes called a bag
258 or multiset. Elements are stored as dictionary keys and their counts
259 are stored as dictionary values.
260
261 >>> c = Counter('abracadabra') # count elements from a string
262
263 >>> c.most_common(3) # three most common elements
264 [('a', 5), ('r', 2), ('b', 2)]
265 >>> sorted(c) # list all unique elements
266 ['a', 'b', 'c', 'd', 'r']
267 >>> ''.join(sorted(c.elements())) # list elements with repetitions
268 'aaaaabbcdrr'
269 >>> sum(c.values()) # total of all counts
270 11
271
272 >>> c['a'] # count of letter 'a'
273 5
274 >>> for elem in 'shazam': # update counts from an iterable
275 ... c[elem] += 1 # by adding 1 to each element's count
276 >>> c['a'] # now there are seven 'a'
277 7
278 >>> del c['r'] # remove all 'r'
279 >>> c['r'] # now there are zero 'r'
280 0
281
282 >>> d = Counter('simsalabim') # make another counter
283 >>> c.update(d) # add in the second counter
284 >>> c['a'] # now there are nine 'a'
285 9
286
287 >>> c.clear() # empty the counter
288 >>> c
289 Counter()
290
291 Note: If a count is set to zero or reduced to zero, it will remain
292 in the counter until the entry is deleted or the counter is cleared:
293
294 >>> c = Counter('aaabbc')
295 >>> c['b'] -= 2 # reduce the count of 'b' by two
296 >>> c.most_common() # 'b' is still in, but its count is zero
297 [('a', 3), ('c', 1), ('b', 0)]
298
299 '''
300 # References:
301 # http://en.wikipedia.org/wiki/Multiset
302 # http://www.gnu.org/software/smalltalk/manual-base/html_node/Bag.html
303 # http://www.demo2s.com/Tutorial/Cpp/0380__set-multiset/Catalog0380__set-multiset.htm
304 # http://code.activestate.com/recipes/259174/
305 # Knuth, TAOCP Vol. II section 4.6.3
306
Raymond Hettinger4d2073a2009-01-20 03:41:22 +0000307 def __init__(self, iterable=None, **kwds):
Raymond Hettingerb8baf632009-01-14 02:20:07 +0000308 '''Create a new, empty Counter object. And if given, count elements
309 from an input iterable. Or, initialize the count from another mapping
310 of elements to their counts.
311
312 >>> c = Counter() # a new, empty counter
313 >>> c = Counter('gallahad') # a new counter from an iterable
314 >>> c = Counter({'a': 4, 'b': 2}) # a new counter from a mapping
Raymond Hettinger4d2073a2009-01-20 03:41:22 +0000315 >>> c = Counter(a=4, b=2) # a new counter from keyword args
Raymond Hettingerb8baf632009-01-14 02:20:07 +0000316
317 '''
Raymond Hettinger4d2073a2009-01-20 03:41:22 +0000318 self.update(iterable, **kwds)
Raymond Hettingerb8baf632009-01-14 02:20:07 +0000319
320 def __missing__(self, key):
321 'The count of elements not in the Counter is zero.'
322 # Needed so that self[missing_item] does not raise KeyError
323 return 0
324
325 def most_common(self, n=None):
326 '''List the n most common elements and their counts from the most
327 common to the least. If n is None, then list all element counts.
328
329 >>> Counter('abracadabra').most_common(3)
330 [('a', 5), ('r', 2), ('b', 2)]
331
332 '''
333 # Emulate Bag.sortedByCount from Smalltalk
334 if n is None:
335 return sorted(self.items(), key=_itemgetter(1), reverse=True)
336 return _heapq.nlargest(n, self.items(), key=_itemgetter(1))
337
338 def elements(self):
339 '''Iterator over elements repeating each as many times as its count.
340
341 >>> c = Counter('ABCABC')
342 >>> sorted(c.elements())
343 ['A', 'A', 'B', 'B', 'C', 'C']
344
345 # Knuth's example for prime factors of 1836: 2**2 * 3**3 * 17**1
346 >>> prime_factors = Counter({2: 2, 3: 3, 17: 1})
347 >>> product = 1
348 >>> for factor in prime_factors.elements(): # loop over factors
349 ... product *= factor # and multiply them
350 >>> product
351 1836
352
353 Note, if an element's count has been set to zero or is a negative
354 number, elements() will ignore it.
355
356 '''
357 # Emulate Bag.do from Smalltalk and Multiset.begin from C++.
358 return _chain.from_iterable(_starmap(_repeat, self.items()))
359
360 # Override dict methods where necessary
361
362 @classmethod
363 def fromkeys(cls, iterable, v=None):
364 # There is no equivalent method for counters because setting v=1
365 # means that no element can have a count greater than one.
366 raise NotImplementedError(
367 'Counter.fromkeys() is undefined. Use Counter(iterable) instead.')
368
Raymond Hettinger4d2073a2009-01-20 03:41:22 +0000369 def update(self, iterable=None, **kwds):
Raymond Hettingerb8baf632009-01-14 02:20:07 +0000370 '''Like dict.update() but add counts instead of replacing them.
371
372 Source can be an iterable, a dictionary, or another Counter instance.
373
374 >>> c = Counter('which')
375 >>> c.update('witch') # add elements from another iterable
376 >>> d = Counter('watch')
377 >>> c.update(d) # add elements from another counter
378 >>> c['h'] # four 'h' in which, witch, and watch
379 4
380
381 '''
382 # The regular dict.update() operation makes no sense here because the
383 # replace behavior results in the some of original untouched counts
384 # being mixed-in with all of the other counts for a mismash that
385 # doesn't have a straight-forward interpretation in most counting
Raymond Hettinger4d2073a2009-01-20 03:41:22 +0000386 # contexts. Instead, we implement straight-addition. Both the inputs
387 # and outputs are allowed to contain zero and negative counts.
Raymond Hettingerb8baf632009-01-14 02:20:07 +0000388
389 if iterable is not None:
390 if isinstance(iterable, Mapping):
Raymond Hettingerdd01f8f2009-01-22 09:09:55 +0000391 if self:
392 for elem, count in iterable.items():
393 self[elem] += count
394 else:
395 dict.update(self, iterable) # fast path when counter is empty
Raymond Hettingerb8baf632009-01-14 02:20:07 +0000396 else:
397 for elem in iterable:
398 self[elem] += 1
Raymond Hettinger4d2073a2009-01-20 03:41:22 +0000399 if kwds:
400 self.update(kwds)
Raymond Hettingerb8baf632009-01-14 02:20:07 +0000401
402 def copy(self):
403 'Like dict.copy() but returns a Counter instance instead of a dict.'
404 return Counter(self)
405
Raymond Hettinger4d2073a2009-01-20 03:41:22 +0000406 def __delitem__(self, elem):
407 'Like dict.__delitem__() but does not raise KeyError for missing values.'
408 if elem in self:
409 dict.__delitem__(self, elem)
410
Raymond Hettingerb8baf632009-01-14 02:20:07 +0000411 def __repr__(self):
412 if not self:
413 return '%s()' % self.__class__.__name__
414 items = ', '.join(map('%r: %r'.__mod__, self.most_common()))
415 return '%s({%s})' % (self.__class__.__name__, items)
416
Raymond Hettinger4d2073a2009-01-20 03:41:22 +0000417 # Multiset-style mathematical operations discussed in:
418 # Knuth TAOCP Volume II section 4.6.3 exercise 19
419 # and at http://en.wikipedia.org/wiki/Multiset
420 #
Raymond Hettinger4d2073a2009-01-20 03:41:22 +0000421 # Outputs guaranteed to only include positive counts.
422 #
423 # To strip negative and zero counts, add-in an empty counter:
424 # c += Counter()
425
426 def __add__(self, other):
427 '''Add counts from two counters.
428
429 >>> Counter('abbb') + Counter('bcc')
430 Counter({'b': 4, 'c': 2, 'a': 1})
431
432 '''
433 if not isinstance(other, Counter):
434 return NotImplemented
435 result = Counter()
436 for elem in set(self) | set(other):
437 newcount = self[elem] + other[elem]
438 if newcount > 0:
439 result[elem] = newcount
440 return result
441
442 def __sub__(self, other):
443 ''' Subtract count, but keep only results with positive counts.
444
445 >>> Counter('abbbc') - Counter('bccd')
446 Counter({'b': 2, 'a': 1})
447
448 '''
449 if not isinstance(other, Counter):
450 return NotImplemented
451 result = Counter()
Raymond Hettingere0d1b9f2009-01-21 20:36:27 +0000452 for elem in set(self) | set(other):
453 newcount = self[elem] - other[elem]
Raymond Hettinger4d2073a2009-01-20 03:41:22 +0000454 if newcount > 0:
455 result[elem] = newcount
456 return result
457
458 def __or__(self, other):
459 '''Union is the maximum of value in either of the input counters.
460
461 >>> Counter('abbb') | Counter('bcc')
462 Counter({'b': 3, 'c': 2, 'a': 1})
463
464 '''
465 if not isinstance(other, Counter):
466 return NotImplemented
Raymond Hettinger4d2073a2009-01-20 03:41:22 +0000467 result = Counter()
468 for elem in set(self) | set(other):
Raymond Hettingerc4791702009-04-04 08:48:03 +0000469 p, q = self[elem], other[elem]
470 newcount = q if p < q else p
Raymond Hettinger4d2073a2009-01-20 03:41:22 +0000471 if newcount > 0:
472 result[elem] = newcount
473 return result
474
475 def __and__(self, other):
476 ''' Intersection is the minimum of corresponding counts.
477
478 >>> Counter('abbb') & Counter('bcc')
479 Counter({'b': 1})
480
481 '''
482 if not isinstance(other, Counter):
483 return NotImplemented
Raymond Hettinger4d2073a2009-01-20 03:41:22 +0000484 result = Counter()
485 if len(self) < len(other):
486 self, other = other, self
487 for elem in filter(self.__contains__, other):
Raymond Hettingerc4791702009-04-04 08:48:03 +0000488 p, q = self[elem], other[elem]
489 newcount = p if p < q else q
Raymond Hettinger4d2073a2009-01-20 03:41:22 +0000490 if newcount > 0:
491 result[elem] = newcount
492 return result
493
Guido van Rossumd8faa362007-04-27 19:54:29 +0000494
Raymond Hettinger48b8b662008-02-05 01:53:00 +0000495################################################################################
496### UserDict
497################################################################################
Guido van Rossumd8faa362007-04-27 19:54:29 +0000498
Raymond Hettinger48b8b662008-02-05 01:53:00 +0000499class UserDict(MutableMapping):
500
501 # Start by filling-out the abstract methods
502 def __init__(self, dict=None, **kwargs):
503 self.data = {}
504 if dict is not None:
505 self.update(dict)
506 if len(kwargs):
507 self.update(kwargs)
508 def __len__(self): return len(self.data)
509 def __getitem__(self, key):
510 if key in self.data:
511 return self.data[key]
512 if hasattr(self.__class__, "__missing__"):
513 return self.__class__.__missing__(self, key)
514 raise KeyError(key)
515 def __setitem__(self, key, item): self.data[key] = item
516 def __delitem__(self, key): del self.data[key]
517 def __iter__(self):
518 return iter(self.data)
519
Raymond Hettinger554c8b82008-02-05 22:54:43 +0000520 # Modify __contains__ to work correctly when __missing__ is present
521 def __contains__(self, key):
522 return key in self.data
Raymond Hettinger48b8b662008-02-05 01:53:00 +0000523
524 # Now, add the methods in dicts but not in MutableMapping
525 def __repr__(self): return repr(self.data)
526 def copy(self):
527 if self.__class__ is UserDict:
528 return UserDict(self.data.copy())
529 import copy
530 data = self.data
531 try:
532 self.data = {}
533 c = copy.copy(self)
534 finally:
535 self.data = data
536 c.update(self)
537 return c
538 @classmethod
539 def fromkeys(cls, iterable, value=None):
540 d = cls()
541 for key in iterable:
542 d[key] = value
543 return d
544
Raymond Hettinger48b8b662008-02-05 01:53:00 +0000545
546
547################################################################################
Raymond Hettinger53dbe392008-02-12 20:03:09 +0000548### UserList
549################################################################################
550
551class UserList(MutableSequence):
552 """A more or less complete user-defined wrapper around list objects."""
553 def __init__(self, initlist=None):
554 self.data = []
555 if initlist is not None:
556 # XXX should this accept an arbitrary sequence?
557 if type(initlist) == type(self.data):
558 self.data[:] = initlist
559 elif isinstance(initlist, UserList):
560 self.data[:] = initlist.data[:]
561 else:
562 self.data = list(initlist)
563 def __repr__(self): return repr(self.data)
564 def __lt__(self, other): return self.data < self.__cast(other)
565 def __le__(self, other): return self.data <= self.__cast(other)
566 def __eq__(self, other): return self.data == self.__cast(other)
567 def __ne__(self, other): return self.data != self.__cast(other)
568 def __gt__(self, other): return self.data > self.__cast(other)
569 def __ge__(self, other): return self.data >= self.__cast(other)
570 def __cast(self, other):
571 return other.data if isinstance(other, UserList) else other
Raymond Hettinger53dbe392008-02-12 20:03:09 +0000572 def __contains__(self, item): return item in self.data
573 def __len__(self): return len(self.data)
574 def __getitem__(self, i): return self.data[i]
575 def __setitem__(self, i, item): self.data[i] = item
576 def __delitem__(self, i): del self.data[i]
577 def __add__(self, other):
578 if isinstance(other, UserList):
579 return self.__class__(self.data + other.data)
580 elif isinstance(other, type(self.data)):
581 return self.__class__(self.data + other)
582 return self.__class__(self.data + list(other))
583 def __radd__(self, other):
584 if isinstance(other, UserList):
585 return self.__class__(other.data + self.data)
586 elif isinstance(other, type(self.data)):
587 return self.__class__(other + self.data)
588 return self.__class__(list(other) + self.data)
589 def __iadd__(self, other):
590 if isinstance(other, UserList):
591 self.data += other.data
592 elif isinstance(other, type(self.data)):
593 self.data += other
594 else:
595 self.data += list(other)
596 return self
597 def __mul__(self, n):
598 return self.__class__(self.data*n)
599 __rmul__ = __mul__
600 def __imul__(self, n):
601 self.data *= n
602 return self
603 def append(self, item): self.data.append(item)
604 def insert(self, i, item): self.data.insert(i, item)
605 def pop(self, i=-1): return self.data.pop(i)
606 def remove(self, item): self.data.remove(item)
607 def count(self, item): return self.data.count(item)
608 def index(self, item, *args): return self.data.index(item, *args)
609 def reverse(self): self.data.reverse()
610 def sort(self, *args, **kwds): self.data.sort(*args, **kwds)
611 def extend(self, other):
612 if isinstance(other, UserList):
613 self.data.extend(other.data)
614 else:
615 self.data.extend(other)
616
617
618
619################################################################################
Raymond Hettingerb3a65f82008-02-21 22:11:37 +0000620### UserString
621################################################################################
622
623class UserString(Sequence):
624 def __init__(self, seq):
625 if isinstance(seq, str):
626 self.data = seq
627 elif isinstance(seq, UserString):
628 self.data = seq.data[:]
629 else:
630 self.data = str(seq)
631 def __str__(self): return str(self.data)
632 def __repr__(self): return repr(self.data)
633 def __int__(self): return int(self.data)
Raymond Hettingerb3a65f82008-02-21 22:11:37 +0000634 def __float__(self): return float(self.data)
635 def __complex__(self): return complex(self.data)
636 def __hash__(self): return hash(self.data)
637
638 def __eq__(self, string):
639 if isinstance(string, UserString):
640 return self.data == string.data
641 return self.data == string
642 def __ne__(self, string):
643 if isinstance(string, UserString):
644 return self.data != string.data
645 return self.data != string
646 def __lt__(self, string):
647 if isinstance(string, UserString):
648 return self.data < string.data
649 return self.data < string
650 def __le__(self, string):
651 if isinstance(string, UserString):
652 return self.data <= string.data
653 return self.data <= string
654 def __gt__(self, string):
655 if isinstance(string, UserString):
656 return self.data > string.data
657 return self.data > string
658 def __ge__(self, string):
659 if isinstance(string, UserString):
660 return self.data >= string.data
661 return self.data >= string
662
663 def __contains__(self, char):
664 if isinstance(char, UserString):
665 char = char.data
666 return char in self.data
667
668 def __len__(self): return len(self.data)
669 def __getitem__(self, index): return self.__class__(self.data[index])
670 def __add__(self, other):
671 if isinstance(other, UserString):
672 return self.__class__(self.data + other.data)
673 elif isinstance(other, str):
674 return self.__class__(self.data + other)
675 return self.__class__(self.data + str(other))
676 def __radd__(self, other):
677 if isinstance(other, str):
678 return self.__class__(other + self.data)
679 return self.__class__(str(other) + self.data)
680 def __mul__(self, n):
681 return self.__class__(self.data*n)
682 __rmul__ = __mul__
683 def __mod__(self, args):
684 return self.__class__(self.data % args)
685
686 # the following methods are defined in alphabetical order:
687 def capitalize(self): return self.__class__(self.data.capitalize())
688 def center(self, width, *args):
689 return self.__class__(self.data.center(width, *args))
690 def count(self, sub, start=0, end=_sys.maxsize):
691 if isinstance(sub, UserString):
692 sub = sub.data
693 return self.data.count(sub, start, end)
694 def encode(self, encoding=None, errors=None): # XXX improve this?
695 if encoding:
696 if errors:
697 return self.__class__(self.data.encode(encoding, errors))
698 return self.__class__(self.data.encode(encoding))
699 return self.__class__(self.data.encode())
700 def endswith(self, suffix, start=0, end=_sys.maxsize):
701 return self.data.endswith(suffix, start, end)
702 def expandtabs(self, tabsize=8):
703 return self.__class__(self.data.expandtabs(tabsize))
704 def find(self, sub, start=0, end=_sys.maxsize):
705 if isinstance(sub, UserString):
706 sub = sub.data
707 return self.data.find(sub, start, end)
708 def format(self, *args, **kwds):
709 return self.data.format(*args, **kwds)
710 def index(self, sub, start=0, end=_sys.maxsize):
711 return self.data.index(sub, start, end)
712 def isalpha(self): return self.data.isalpha()
713 def isalnum(self): return self.data.isalnum()
714 def isdecimal(self): return self.data.isdecimal()
715 def isdigit(self): return self.data.isdigit()
716 def isidentifier(self): return self.data.isidentifier()
717 def islower(self): return self.data.islower()
718 def isnumeric(self): return self.data.isnumeric()
719 def isspace(self): return self.data.isspace()
720 def istitle(self): return self.data.istitle()
721 def isupper(self): return self.data.isupper()
722 def join(self, seq): return self.data.join(seq)
723 def ljust(self, width, *args):
724 return self.__class__(self.data.ljust(width, *args))
725 def lower(self): return self.__class__(self.data.lower())
726 def lstrip(self, chars=None): return self.__class__(self.data.lstrip(chars))
727 def partition(self, sep):
728 return self.data.partition(sep)
729 def replace(self, old, new, maxsplit=-1):
730 if isinstance(old, UserString):
731 old = old.data
732 if isinstance(new, UserString):
733 new = new.data
734 return self.__class__(self.data.replace(old, new, maxsplit))
735 def rfind(self, sub, start=0, end=_sys.maxsize):
736 return self.data.rfind(sub, start, end)
737 def rindex(self, sub, start=0, end=_sys.maxsize):
738 return self.data.rindex(sub, start, end)
739 def rjust(self, width, *args):
740 return self.__class__(self.data.rjust(width, *args))
741 def rpartition(self, sep):
742 return self.data.rpartition(sep)
743 def rstrip(self, chars=None):
744 return self.__class__(self.data.rstrip(chars))
745 def split(self, sep=None, maxsplit=-1):
746 return self.data.split(sep, maxsplit)
747 def rsplit(self, sep=None, maxsplit=-1):
748 return self.data.rsplit(sep, maxsplit)
749 def splitlines(self, keepends=0): return self.data.splitlines(keepends)
750 def startswith(self, prefix, start=0, end=_sys.maxsize):
751 return self.data.startswith(prefix, start, end)
752 def strip(self, chars=None): return self.__class__(self.data.strip(chars))
753 def swapcase(self): return self.__class__(self.data.swapcase())
754 def title(self): return self.__class__(self.data.title())
755 def translate(self, *args):
756 return self.__class__(self.data.translate(*args))
757 def upper(self): return self.__class__(self.data.upper())
758 def zfill(self, width): return self.__class__(self.data.zfill(width))
759
760
761
762################################################################################
Raymond Hettinger48b8b662008-02-05 01:53:00 +0000763### Simple tests
764################################################################################
Guido van Rossumd8faa362007-04-27 19:54:29 +0000765
766if __name__ == '__main__':
Thomas Wouters1b7f8912007-09-19 03:06:30 +0000767 # verify that instances can be pickled
Guido van Rossum99603b02007-07-20 00:22:32 +0000768 from pickle import loads, dumps
Guido van Rossum8ce8a782007-11-01 19:42:39 +0000769 Point = namedtuple('Point', 'x, y', True)
Guido van Rossumd8faa362007-04-27 19:54:29 +0000770 p = Point(x=10, y=20)
771 assert p == loads(dumps(p))
772
Guido van Rossum3d392eb2007-11-16 00:35:22 +0000773 # test and demonstrate ability to override methods
Christian Heimes043d6f62008-01-07 17:19:16 +0000774 class Point(namedtuple('Point', 'x y')):
Christian Heimes25bb7832008-01-11 16:17:00 +0000775 __slots__ = ()
Christian Heimes043d6f62008-01-07 17:19:16 +0000776 @property
777 def hypot(self):
778 return (self.x ** 2 + self.y ** 2) ** 0.5
Christian Heimes790c8232008-01-07 21:14:23 +0000779 def __str__(self):
Christian Heimes25bb7832008-01-11 16:17:00 +0000780 return 'Point: x=%6.3f y=%6.3f hypot=%6.3f' % (self.x, self.y, self.hypot)
Christian Heimes043d6f62008-01-07 17:19:16 +0000781
Christian Heimes25bb7832008-01-11 16:17:00 +0000782 for p in Point(3, 4), Point(14, 5/7.):
Christian Heimes790c8232008-01-07 21:14:23 +0000783 print (p)
Christian Heimes043d6f62008-01-07 17:19:16 +0000784
785 class Point(namedtuple('Point', 'x y')):
786 'Point class with optimized _make() and _replace() without error-checking'
Christian Heimes25bb7832008-01-11 16:17:00 +0000787 __slots__ = ()
Christian Heimes043d6f62008-01-07 17:19:16 +0000788 _make = classmethod(tuple.__new__)
789 def _replace(self, _map=map, **kwds):
Christian Heimes2380ac72008-01-09 00:17:24 +0000790 return self._make(_map(kwds.get, ('x', 'y'), self))
Christian Heimes043d6f62008-01-07 17:19:16 +0000791
792 print(Point(11, 22)._replace(x=100))
Guido van Rossum3d392eb2007-11-16 00:35:22 +0000793
Christian Heimes25bb7832008-01-11 16:17:00 +0000794 Point3D = namedtuple('Point3D', Point._fields + ('z',))
795 print(Point3D.__doc__)
796
Guido van Rossumd8faa362007-04-27 19:54:29 +0000797 import doctest
Guido van Rossum8ce8a782007-11-01 19:42:39 +0000798 TestResults = namedtuple('TestResults', 'failed attempted')
Guido van Rossumd8faa362007-04-27 19:54:29 +0000799 print(TestResults(*doctest.testmod()))