blob: 002ce9724de929dc9437f81035a07f456ca7b978 [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
Raymond Hettingerb8baf632009-01-14 02:20:07 +0000135
Christian Heimes99170a52007-12-19 02:07:34 +0000136
Raymond Hettinger48b8b662008-02-05 01:53:00 +0000137################################################################################
138### namedtuple
139################################################################################
140
Benjamin Petersona86f2c02009-02-10 02:41:10 +0000141def namedtuple(typename, field_names, verbose=False, rename=False):
Guido van Rossumd8faa362007-04-27 19:54:29 +0000142 """Returns a new subclass of tuple with named fields.
143
Guido van Rossum8ce8a782007-11-01 19:42:39 +0000144 >>> Point = namedtuple('Point', 'x y')
Thomas Wouters1b7f8912007-09-19 03:06:30 +0000145 >>> Point.__doc__ # docstring for the new class
Guido van Rossumd8faa362007-04-27 19:54:29 +0000146 'Point(x, y)'
Thomas Wouters1b7f8912007-09-19 03:06:30 +0000147 >>> p = Point(11, y=22) # instantiate with positional args or keywords
Christian Heimes99170a52007-12-19 02:07:34 +0000148 >>> p[0] + p[1] # indexable like a plain tuple
Guido van Rossumd8faa362007-04-27 19:54:29 +0000149 33
Christian Heimes99170a52007-12-19 02:07:34 +0000150 >>> x, y = p # unpack like a regular tuple
Guido van Rossumd8faa362007-04-27 19:54:29 +0000151 >>> x, y
152 (11, 22)
Thomas Wouters1b7f8912007-09-19 03:06:30 +0000153 >>> p.x + p.y # fields also accessable by name
Guido van Rossumd8faa362007-04-27 19:54:29 +0000154 33
Christian Heimes0449f632007-12-15 01:27:15 +0000155 >>> d = p._asdict() # convert to a dictionary
Guido van Rossum8ce8a782007-11-01 19:42:39 +0000156 >>> d['x']
157 11
158 >>> Point(**d) # convert from a dictionary
Guido van Rossumd8faa362007-04-27 19:54:29 +0000159 Point(x=11, y=22)
Christian Heimes0449f632007-12-15 01:27:15 +0000160 >>> p._replace(x=100) # _replace() is like str.replace() but targets named fields
Thomas Wouters1b7f8912007-09-19 03:06:30 +0000161 Point(x=100, y=22)
Guido van Rossumd8faa362007-04-27 19:54:29 +0000162
163 """
164
Christian Heimes2380ac72008-01-09 00:17:24 +0000165 # Parse and validate the field names. Validation serves two purposes,
166 # generating informative error messages and preventing template injection attacks.
Guido van Rossum8ce8a782007-11-01 19:42:39 +0000167 if isinstance(field_names, str):
168 field_names = field_names.replace(',', ' ').split() # names separated by whitespace and/or commas
Benjamin Petersone9bbc8b2008-09-28 02:06:32 +0000169 field_names = tuple(map(str, field_names))
Benjamin Petersona86f2c02009-02-10 02:41:10 +0000170 if rename:
171 names = list(field_names)
172 seen = set()
173 for i, name in enumerate(names):
174 if (not all(c.isalnum() or c=='_' for c in name) or _iskeyword(name)
175 or not name or name[0].isdigit() or name.startswith('_')
176 or name in seen):
177 names[i] = '_%d' % (i+1)
178 seen.add(name)
179 field_names = tuple(names)
Guido van Rossum8ce8a782007-11-01 19:42:39 +0000180 for name in (typename,) + field_names:
Christian Heimesb9eccbf2007-12-05 20:18:38 +0000181 if not all(c.isalnum() or c=='_' for c in name):
Guido van Rossum8ce8a782007-11-01 19:42:39 +0000182 raise ValueError('Type names and field names can only contain alphanumeric characters and underscores: %r' % name)
183 if _iskeyword(name):
184 raise ValueError('Type names and field names cannot be a keyword: %r' % name)
185 if name[0].isdigit():
186 raise ValueError('Type names and field names cannot start with a number: %r' % name)
187 seen_names = set()
188 for name in field_names:
Benjamin Petersona86f2c02009-02-10 02:41:10 +0000189 if name.startswith('_') and not rename:
Christian Heimes0449f632007-12-15 01:27:15 +0000190 raise ValueError('Field names cannot start with an underscore: %r' % name)
Guido van Rossum8ce8a782007-11-01 19:42:39 +0000191 if name in seen_names:
192 raise ValueError('Encountered duplicate field name: %r' % name)
193 seen_names.add(name)
194
195 # Create and fill-in the class template
Christian Heimesfaf2f632008-01-06 16:59:19 +0000196 numfields = len(field_names)
Thomas Wouters1b7f8912007-09-19 03:06:30 +0000197 argtxt = repr(field_names).replace("'", "")[1:-1] # tuple repr without parens or quotes
Guido van Rossumd59da4b2007-05-22 18:11:13 +0000198 reprtxt = ', '.join('%s=%%r' % name for name in field_names)
199 template = '''class %(typename)s(tuple):
Christian Heimes0449f632007-12-15 01:27:15 +0000200 '%(typename)s(%(argtxt)s)' \n
201 __slots__ = () \n
Christian Heimesfaf2f632008-01-06 16:59:19 +0000202 _fields = %(field_names)r \n
Guido van Rossumd59da4b2007-05-22 18:11:13 +0000203 def __new__(cls, %(argtxt)s):
Christian Heimes0449f632007-12-15 01:27:15 +0000204 return tuple.__new__(cls, (%(argtxt)s)) \n
Christian Heimesfaf2f632008-01-06 16:59:19 +0000205 @classmethod
Christian Heimes043d6f62008-01-07 17:19:16 +0000206 def _make(cls, iterable, new=tuple.__new__, len=len):
Christian Heimesfaf2f632008-01-06 16:59:19 +0000207 'Make a new %(typename)s object from a sequence or iterable'
Christian Heimes043d6f62008-01-07 17:19:16 +0000208 result = new(cls, iterable)
Christian Heimesfaf2f632008-01-06 16:59:19 +0000209 if len(result) != %(numfields)d:
210 raise TypeError('Expected %(numfields)d arguments, got %%d' %% len(result))
211 return result \n
Guido van Rossumd59da4b2007-05-22 18:11:13 +0000212 def __repr__(self):
Christian Heimes0449f632007-12-15 01:27:15 +0000213 return '%(typename)s(%(reprtxt)s)' %% self \n
Raymond Hettingera4f52b12009-03-02 22:28:31 +0000214 def _asdict(self):
215 'Return a new OrderedDict which maps field names to their values'
216 return OrderedDict(zip(self._fields, self)) \n
Christian Heimes0449f632007-12-15 01:27:15 +0000217 def _replace(self, **kwds):
Guido van Rossum3d392eb2007-11-16 00:35:22 +0000218 'Return a new %(typename)s object replacing specified fields with new values'
Christian Heimesfaf2f632008-01-06 16:59:19 +0000219 result = self._make(map(kwds.pop, %(field_names)r, self))
220 if kwds:
221 raise ValueError('Got unexpected field names: %%r' %% kwds.keys())
Georg Brandlc28e1fa2008-06-10 19:20:26 +0000222 return result \n
223 def __getnewargs__(self):
224 return tuple(self) \n\n''' % locals()
Guido van Rossumd59da4b2007-05-22 18:11:13 +0000225 for i, name in enumerate(field_names):
Thomas Wouters1b7f8912007-09-19 03:06:30 +0000226 template += ' %s = property(itemgetter(%d))\n' % (name, i)
227 if verbose:
228 print(template)
Guido van Rossum8ce8a782007-11-01 19:42:39 +0000229
Georg Brandlf08a9dd2008-06-10 16:57:31 +0000230 # Execute the template string in a temporary namespace and
231 # support tracing utilities by setting a value for frame.f_globals['__name__']
Raymond Hettingera4f52b12009-03-02 22:28:31 +0000232 namespace = dict(itemgetter=_itemgetter, __name__='namedtuple_%s' % typename,
233 OrderedDict=OrderedDict)
Guido van Rossum8ce8a782007-11-01 19:42:39 +0000234 try:
235 exec(template, namespace)
236 except SyntaxError as e:
Christian Heimes99170a52007-12-19 02:07:34 +0000237 raise SyntaxError(e.msg + ':\n' + template) from e
Guido van Rossum8ce8a782007-11-01 19:42:39 +0000238 result = namespace[typename]
239
240 # For pickling to work, the __module__ variable needs to be set to the frame
241 # where the named tuple is created. Bypass this step in enviroments where
242 # sys._getframe is not defined (Jython for example).
Guido van Rossumd59da4b2007-05-22 18:11:13 +0000243 if hasattr(_sys, '_getframe'):
Raymond Hettinger0f055172009-01-27 10:06:09 +0000244 result.__module__ = _sys._getframe(1).f_globals.get('__name__', '__main__')
Guido van Rossum8ce8a782007-11-01 19:42:39 +0000245
Guido van Rossumd59da4b2007-05-22 18:11:13 +0000246 return result
Guido van Rossumd8faa362007-04-27 19:54:29 +0000247
Guido van Rossumd8faa362007-04-27 19:54:29 +0000248
Raymond Hettingerb8baf632009-01-14 02:20:07 +0000249########################################################################
250### Counter
251########################################################################
252
253class Counter(dict):
254 '''Dict subclass for counting hashable items. Sometimes called a bag
255 or multiset. Elements are stored as dictionary keys and their counts
256 are stored as dictionary values.
257
258 >>> c = Counter('abracadabra') # count elements from a string
259
260 >>> c.most_common(3) # three most common elements
261 [('a', 5), ('r', 2), ('b', 2)]
262 >>> sorted(c) # list all unique elements
263 ['a', 'b', 'c', 'd', 'r']
264 >>> ''.join(sorted(c.elements())) # list elements with repetitions
265 'aaaaabbcdrr'
266 >>> sum(c.values()) # total of all counts
267 11
268
269 >>> c['a'] # count of letter 'a'
270 5
271 >>> for elem in 'shazam': # update counts from an iterable
272 ... c[elem] += 1 # by adding 1 to each element's count
273 >>> c['a'] # now there are seven 'a'
274 7
275 >>> del c['r'] # remove all 'r'
276 >>> c['r'] # now there are zero 'r'
277 0
278
279 >>> d = Counter('simsalabim') # make another counter
280 >>> c.update(d) # add in the second counter
281 >>> c['a'] # now there are nine 'a'
282 9
283
284 >>> c.clear() # empty the counter
285 >>> c
286 Counter()
287
288 Note: If a count is set to zero or reduced to zero, it will remain
289 in the counter until the entry is deleted or the counter is cleared:
290
291 >>> c = Counter('aaabbc')
292 >>> c['b'] -= 2 # reduce the count of 'b' by two
293 >>> c.most_common() # 'b' is still in, but its count is zero
294 [('a', 3), ('c', 1), ('b', 0)]
295
296 '''
297 # References:
298 # http://en.wikipedia.org/wiki/Multiset
299 # http://www.gnu.org/software/smalltalk/manual-base/html_node/Bag.html
300 # http://www.demo2s.com/Tutorial/Cpp/0380__set-multiset/Catalog0380__set-multiset.htm
301 # http://code.activestate.com/recipes/259174/
302 # Knuth, TAOCP Vol. II section 4.6.3
303
Raymond Hettinger4d2073a2009-01-20 03:41:22 +0000304 def __init__(self, iterable=None, **kwds):
Raymond Hettingerb8baf632009-01-14 02:20:07 +0000305 '''Create a new, empty Counter object. And if given, count elements
306 from an input iterable. Or, initialize the count from another mapping
307 of elements to their counts.
308
309 >>> c = Counter() # a new, empty counter
310 >>> c = Counter('gallahad') # a new counter from an iterable
311 >>> c = Counter({'a': 4, 'b': 2}) # a new counter from a mapping
Raymond Hettinger4d2073a2009-01-20 03:41:22 +0000312 >>> c = Counter(a=4, b=2) # a new counter from keyword args
Raymond Hettingerb8baf632009-01-14 02:20:07 +0000313
314 '''
Raymond Hettinger4d2073a2009-01-20 03:41:22 +0000315 self.update(iterable, **kwds)
Raymond Hettingerb8baf632009-01-14 02:20:07 +0000316
317 def __missing__(self, key):
318 'The count of elements not in the Counter is zero.'
319 # Needed so that self[missing_item] does not raise KeyError
320 return 0
321
322 def most_common(self, n=None):
323 '''List the n most common elements and their counts from the most
324 common to the least. If n is None, then list all element counts.
325
326 >>> Counter('abracadabra').most_common(3)
327 [('a', 5), ('r', 2), ('b', 2)]
328
329 '''
330 # Emulate Bag.sortedByCount from Smalltalk
331 if n is None:
332 return sorted(self.items(), key=_itemgetter(1), reverse=True)
333 return _heapq.nlargest(n, self.items(), key=_itemgetter(1))
334
335 def elements(self):
336 '''Iterator over elements repeating each as many times as its count.
337
338 >>> c = Counter('ABCABC')
339 >>> sorted(c.elements())
340 ['A', 'A', 'B', 'B', 'C', 'C']
341
342 # Knuth's example for prime factors of 1836: 2**2 * 3**3 * 17**1
343 >>> prime_factors = Counter({2: 2, 3: 3, 17: 1})
344 >>> product = 1
345 >>> for factor in prime_factors.elements(): # loop over factors
346 ... product *= factor # and multiply them
347 >>> product
348 1836
349
350 Note, if an element's count has been set to zero or is a negative
351 number, elements() will ignore it.
352
353 '''
354 # Emulate Bag.do from Smalltalk and Multiset.begin from C++.
355 return _chain.from_iterable(_starmap(_repeat, self.items()))
356
357 # Override dict methods where necessary
358
359 @classmethod
360 def fromkeys(cls, iterable, v=None):
361 # There is no equivalent method for counters because setting v=1
362 # means that no element can have a count greater than one.
363 raise NotImplementedError(
364 'Counter.fromkeys() is undefined. Use Counter(iterable) instead.')
365
Raymond Hettinger4d2073a2009-01-20 03:41:22 +0000366 def update(self, iterable=None, **kwds):
Raymond Hettingerb8baf632009-01-14 02:20:07 +0000367 '''Like dict.update() but add counts instead of replacing them.
368
369 Source can be an iterable, a dictionary, or another Counter instance.
370
371 >>> c = Counter('which')
372 >>> c.update('witch') # add elements from another iterable
373 >>> d = Counter('watch')
374 >>> c.update(d) # add elements from another counter
375 >>> c['h'] # four 'h' in which, witch, and watch
376 4
377
378 '''
379 # The regular dict.update() operation makes no sense here because the
380 # replace behavior results in the some of original untouched counts
381 # being mixed-in with all of the other counts for a mismash that
382 # doesn't have a straight-forward interpretation in most counting
Raymond Hettinger4d2073a2009-01-20 03:41:22 +0000383 # contexts. Instead, we implement straight-addition. Both the inputs
384 # and outputs are allowed to contain zero and negative counts.
Raymond Hettingerb8baf632009-01-14 02:20:07 +0000385
386 if iterable is not None:
387 if isinstance(iterable, Mapping):
Raymond Hettingerdd01f8f2009-01-22 09:09:55 +0000388 if self:
389 for elem, count in iterable.items():
390 self[elem] += count
391 else:
392 dict.update(self, iterable) # fast path when counter is empty
Raymond Hettingerb8baf632009-01-14 02:20:07 +0000393 else:
394 for elem in iterable:
395 self[elem] += 1
Raymond Hettinger4d2073a2009-01-20 03:41:22 +0000396 if kwds:
397 self.update(kwds)
Raymond Hettingerb8baf632009-01-14 02:20:07 +0000398
399 def copy(self):
400 'Like dict.copy() but returns a Counter instance instead of a dict.'
401 return Counter(self)
402
Raymond Hettinger4d2073a2009-01-20 03:41:22 +0000403 def __delitem__(self, elem):
404 'Like dict.__delitem__() but does not raise KeyError for missing values.'
405 if elem in self:
406 dict.__delitem__(self, elem)
407
Raymond Hettingerb8baf632009-01-14 02:20:07 +0000408 def __repr__(self):
409 if not self:
410 return '%s()' % self.__class__.__name__
411 items = ', '.join(map('%r: %r'.__mod__, self.most_common()))
412 return '%s({%s})' % (self.__class__.__name__, items)
413
Raymond Hettinger4d2073a2009-01-20 03:41:22 +0000414 # Multiset-style mathematical operations discussed in:
415 # Knuth TAOCP Volume II section 4.6.3 exercise 19
416 # and at http://en.wikipedia.org/wiki/Multiset
417 #
Raymond Hettinger4d2073a2009-01-20 03:41:22 +0000418 # Outputs guaranteed to only include positive counts.
419 #
420 # To strip negative and zero counts, add-in an empty counter:
421 # c += Counter()
422
423 def __add__(self, other):
424 '''Add counts from two counters.
425
426 >>> Counter('abbb') + Counter('bcc')
427 Counter({'b': 4, 'c': 2, 'a': 1})
428
429 '''
430 if not isinstance(other, Counter):
431 return NotImplemented
432 result = Counter()
433 for elem in set(self) | set(other):
434 newcount = self[elem] + other[elem]
435 if newcount > 0:
436 result[elem] = newcount
437 return result
438
439 def __sub__(self, other):
440 ''' Subtract count, but keep only results with positive counts.
441
442 >>> Counter('abbbc') - Counter('bccd')
443 Counter({'b': 2, 'a': 1})
444
445 '''
446 if not isinstance(other, Counter):
447 return NotImplemented
448 result = Counter()
Raymond Hettingere0d1b9f2009-01-21 20:36:27 +0000449 for elem in set(self) | set(other):
450 newcount = self[elem] - other[elem]
Raymond Hettinger4d2073a2009-01-20 03:41:22 +0000451 if newcount > 0:
452 result[elem] = newcount
453 return result
454
455 def __or__(self, other):
456 '''Union is the maximum of value in either of the input counters.
457
458 >>> Counter('abbb') | Counter('bcc')
459 Counter({'b': 3, 'c': 2, 'a': 1})
460
461 '''
462 if not isinstance(other, Counter):
463 return NotImplemented
464 _max = max
465 result = Counter()
466 for elem in set(self) | set(other):
467 newcount = _max(self[elem], other[elem])
468 if newcount > 0:
469 result[elem] = newcount
470 return result
471
472 def __and__(self, other):
473 ''' Intersection is the minimum of corresponding counts.
474
475 >>> Counter('abbb') & Counter('bcc')
476 Counter({'b': 1})
477
478 '''
479 if not isinstance(other, Counter):
480 return NotImplemented
481 _min = min
482 result = Counter()
483 if len(self) < len(other):
484 self, other = other, self
485 for elem in filter(self.__contains__, other):
486 newcount = _min(self[elem], other[elem])
487 if newcount > 0:
488 result[elem] = newcount
489 return result
490
Guido van Rossumd8faa362007-04-27 19:54:29 +0000491
Raymond Hettinger48b8b662008-02-05 01:53:00 +0000492################################################################################
493### UserDict
494################################################################################
Guido van Rossumd8faa362007-04-27 19:54:29 +0000495
Raymond Hettinger48b8b662008-02-05 01:53:00 +0000496class UserDict(MutableMapping):
497
498 # Start by filling-out the abstract methods
499 def __init__(self, dict=None, **kwargs):
500 self.data = {}
501 if dict is not None:
502 self.update(dict)
503 if len(kwargs):
504 self.update(kwargs)
505 def __len__(self): return len(self.data)
506 def __getitem__(self, key):
507 if key in self.data:
508 return self.data[key]
509 if hasattr(self.__class__, "__missing__"):
510 return self.__class__.__missing__(self, key)
511 raise KeyError(key)
512 def __setitem__(self, key, item): self.data[key] = item
513 def __delitem__(self, key): del self.data[key]
514 def __iter__(self):
515 return iter(self.data)
516
Raymond Hettinger554c8b82008-02-05 22:54:43 +0000517 # Modify __contains__ to work correctly when __missing__ is present
518 def __contains__(self, key):
519 return key in self.data
Raymond Hettinger48b8b662008-02-05 01:53:00 +0000520
521 # Now, add the methods in dicts but not in MutableMapping
522 def __repr__(self): return repr(self.data)
523 def copy(self):
524 if self.__class__ is UserDict:
525 return UserDict(self.data.copy())
526 import copy
527 data = self.data
528 try:
529 self.data = {}
530 c = copy.copy(self)
531 finally:
532 self.data = data
533 c.update(self)
534 return c
535 @classmethod
536 def fromkeys(cls, iterable, value=None):
537 d = cls()
538 for key in iterable:
539 d[key] = value
540 return d
541
Raymond Hettinger48b8b662008-02-05 01:53:00 +0000542
543
544################################################################################
Raymond Hettinger53dbe392008-02-12 20:03:09 +0000545### UserList
546################################################################################
547
548class UserList(MutableSequence):
549 """A more or less complete user-defined wrapper around list objects."""
550 def __init__(self, initlist=None):
551 self.data = []
552 if initlist is not None:
553 # XXX should this accept an arbitrary sequence?
554 if type(initlist) == type(self.data):
555 self.data[:] = initlist
556 elif isinstance(initlist, UserList):
557 self.data[:] = initlist.data[:]
558 else:
559 self.data = list(initlist)
560 def __repr__(self): return repr(self.data)
561 def __lt__(self, other): return self.data < self.__cast(other)
562 def __le__(self, other): return self.data <= self.__cast(other)
563 def __eq__(self, other): return self.data == self.__cast(other)
564 def __ne__(self, other): return self.data != self.__cast(other)
565 def __gt__(self, other): return self.data > self.__cast(other)
566 def __ge__(self, other): return self.data >= self.__cast(other)
567 def __cast(self, other):
568 return other.data if isinstance(other, UserList) else other
Raymond Hettinger53dbe392008-02-12 20:03:09 +0000569 def __contains__(self, item): return item in self.data
570 def __len__(self): return len(self.data)
571 def __getitem__(self, i): return self.data[i]
572 def __setitem__(self, i, item): self.data[i] = item
573 def __delitem__(self, i): del self.data[i]
574 def __add__(self, other):
575 if isinstance(other, UserList):
576 return self.__class__(self.data + other.data)
577 elif isinstance(other, type(self.data)):
578 return self.__class__(self.data + other)
579 return self.__class__(self.data + list(other))
580 def __radd__(self, other):
581 if isinstance(other, UserList):
582 return self.__class__(other.data + self.data)
583 elif isinstance(other, type(self.data)):
584 return self.__class__(other + self.data)
585 return self.__class__(list(other) + self.data)
586 def __iadd__(self, other):
587 if isinstance(other, UserList):
588 self.data += other.data
589 elif isinstance(other, type(self.data)):
590 self.data += other
591 else:
592 self.data += list(other)
593 return self
594 def __mul__(self, n):
595 return self.__class__(self.data*n)
596 __rmul__ = __mul__
597 def __imul__(self, n):
598 self.data *= n
599 return self
600 def append(self, item): self.data.append(item)
601 def insert(self, i, item): self.data.insert(i, item)
602 def pop(self, i=-1): return self.data.pop(i)
603 def remove(self, item): self.data.remove(item)
604 def count(self, item): return self.data.count(item)
605 def index(self, item, *args): return self.data.index(item, *args)
606 def reverse(self): self.data.reverse()
607 def sort(self, *args, **kwds): self.data.sort(*args, **kwds)
608 def extend(self, other):
609 if isinstance(other, UserList):
610 self.data.extend(other.data)
611 else:
612 self.data.extend(other)
613
614
615
616################################################################################
Raymond Hettingerb3a65f82008-02-21 22:11:37 +0000617### UserString
618################################################################################
619
620class UserString(Sequence):
621 def __init__(self, seq):
622 if isinstance(seq, str):
623 self.data = seq
624 elif isinstance(seq, UserString):
625 self.data = seq.data[:]
626 else:
627 self.data = str(seq)
628 def __str__(self): return str(self.data)
629 def __repr__(self): return repr(self.data)
630 def __int__(self): return int(self.data)
Raymond Hettingerb3a65f82008-02-21 22:11:37 +0000631 def __float__(self): return float(self.data)
632 def __complex__(self): return complex(self.data)
633 def __hash__(self): return hash(self.data)
634
635 def __eq__(self, string):
636 if isinstance(string, UserString):
637 return self.data == string.data
638 return self.data == string
639 def __ne__(self, string):
640 if isinstance(string, UserString):
641 return self.data != string.data
642 return self.data != string
643 def __lt__(self, string):
644 if isinstance(string, UserString):
645 return self.data < string.data
646 return self.data < string
647 def __le__(self, string):
648 if isinstance(string, UserString):
649 return self.data <= string.data
650 return self.data <= string
651 def __gt__(self, string):
652 if isinstance(string, UserString):
653 return self.data > string.data
654 return self.data > string
655 def __ge__(self, string):
656 if isinstance(string, UserString):
657 return self.data >= string.data
658 return self.data >= string
659
660 def __contains__(self, char):
661 if isinstance(char, UserString):
662 char = char.data
663 return char in self.data
664
665 def __len__(self): return len(self.data)
666 def __getitem__(self, index): return self.__class__(self.data[index])
667 def __add__(self, other):
668 if isinstance(other, UserString):
669 return self.__class__(self.data + other.data)
670 elif isinstance(other, str):
671 return self.__class__(self.data + other)
672 return self.__class__(self.data + str(other))
673 def __radd__(self, other):
674 if isinstance(other, str):
675 return self.__class__(other + self.data)
676 return self.__class__(str(other) + self.data)
677 def __mul__(self, n):
678 return self.__class__(self.data*n)
679 __rmul__ = __mul__
680 def __mod__(self, args):
681 return self.__class__(self.data % args)
682
683 # the following methods are defined in alphabetical order:
684 def capitalize(self): return self.__class__(self.data.capitalize())
685 def center(self, width, *args):
686 return self.__class__(self.data.center(width, *args))
687 def count(self, sub, start=0, end=_sys.maxsize):
688 if isinstance(sub, UserString):
689 sub = sub.data
690 return self.data.count(sub, start, end)
691 def encode(self, encoding=None, errors=None): # XXX improve this?
692 if encoding:
693 if errors:
694 return self.__class__(self.data.encode(encoding, errors))
695 return self.__class__(self.data.encode(encoding))
696 return self.__class__(self.data.encode())
697 def endswith(self, suffix, start=0, end=_sys.maxsize):
698 return self.data.endswith(suffix, start, end)
699 def expandtabs(self, tabsize=8):
700 return self.__class__(self.data.expandtabs(tabsize))
701 def find(self, sub, start=0, end=_sys.maxsize):
702 if isinstance(sub, UserString):
703 sub = sub.data
704 return self.data.find(sub, start, end)
705 def format(self, *args, **kwds):
706 return self.data.format(*args, **kwds)
707 def index(self, sub, start=0, end=_sys.maxsize):
708 return self.data.index(sub, start, end)
709 def isalpha(self): return self.data.isalpha()
710 def isalnum(self): return self.data.isalnum()
711 def isdecimal(self): return self.data.isdecimal()
712 def isdigit(self): return self.data.isdigit()
713 def isidentifier(self): return self.data.isidentifier()
714 def islower(self): return self.data.islower()
715 def isnumeric(self): return self.data.isnumeric()
716 def isspace(self): return self.data.isspace()
717 def istitle(self): return self.data.istitle()
718 def isupper(self): return self.data.isupper()
719 def join(self, seq): return self.data.join(seq)
720 def ljust(self, width, *args):
721 return self.__class__(self.data.ljust(width, *args))
722 def lower(self): return self.__class__(self.data.lower())
723 def lstrip(self, chars=None): return self.__class__(self.data.lstrip(chars))
724 def partition(self, sep):
725 return self.data.partition(sep)
726 def replace(self, old, new, maxsplit=-1):
727 if isinstance(old, UserString):
728 old = old.data
729 if isinstance(new, UserString):
730 new = new.data
731 return self.__class__(self.data.replace(old, new, maxsplit))
732 def rfind(self, sub, start=0, end=_sys.maxsize):
733 return self.data.rfind(sub, start, end)
734 def rindex(self, sub, start=0, end=_sys.maxsize):
735 return self.data.rindex(sub, start, end)
736 def rjust(self, width, *args):
737 return self.__class__(self.data.rjust(width, *args))
738 def rpartition(self, sep):
739 return self.data.rpartition(sep)
740 def rstrip(self, chars=None):
741 return self.__class__(self.data.rstrip(chars))
742 def split(self, sep=None, maxsplit=-1):
743 return self.data.split(sep, maxsplit)
744 def rsplit(self, sep=None, maxsplit=-1):
745 return self.data.rsplit(sep, maxsplit)
746 def splitlines(self, keepends=0): return self.data.splitlines(keepends)
747 def startswith(self, prefix, start=0, end=_sys.maxsize):
748 return self.data.startswith(prefix, start, end)
749 def strip(self, chars=None): return self.__class__(self.data.strip(chars))
750 def swapcase(self): return self.__class__(self.data.swapcase())
751 def title(self): return self.__class__(self.data.title())
752 def translate(self, *args):
753 return self.__class__(self.data.translate(*args))
754 def upper(self): return self.__class__(self.data.upper())
755 def zfill(self, width): return self.__class__(self.data.zfill(width))
756
757
758
759################################################################################
Raymond Hettinger48b8b662008-02-05 01:53:00 +0000760### Simple tests
761################################################################################
Guido van Rossumd8faa362007-04-27 19:54:29 +0000762
763if __name__ == '__main__':
Thomas Wouters1b7f8912007-09-19 03:06:30 +0000764 # verify that instances can be pickled
Guido van Rossum99603b02007-07-20 00:22:32 +0000765 from pickle import loads, dumps
Guido van Rossum8ce8a782007-11-01 19:42:39 +0000766 Point = namedtuple('Point', 'x, y', True)
Guido van Rossumd8faa362007-04-27 19:54:29 +0000767 p = Point(x=10, y=20)
768 assert p == loads(dumps(p))
769
Guido van Rossum3d392eb2007-11-16 00:35:22 +0000770 # test and demonstrate ability to override methods
Christian Heimes043d6f62008-01-07 17:19:16 +0000771 class Point(namedtuple('Point', 'x y')):
Christian Heimes25bb7832008-01-11 16:17:00 +0000772 __slots__ = ()
Christian Heimes043d6f62008-01-07 17:19:16 +0000773 @property
774 def hypot(self):
775 return (self.x ** 2 + self.y ** 2) ** 0.5
Christian Heimes790c8232008-01-07 21:14:23 +0000776 def __str__(self):
Christian Heimes25bb7832008-01-11 16:17:00 +0000777 return 'Point: x=%6.3f y=%6.3f hypot=%6.3f' % (self.x, self.y, self.hypot)
Christian Heimes043d6f62008-01-07 17:19:16 +0000778
Christian Heimes25bb7832008-01-11 16:17:00 +0000779 for p in Point(3, 4), Point(14, 5/7.):
Christian Heimes790c8232008-01-07 21:14:23 +0000780 print (p)
Christian Heimes043d6f62008-01-07 17:19:16 +0000781
782 class Point(namedtuple('Point', 'x y')):
783 'Point class with optimized _make() and _replace() without error-checking'
Christian Heimes25bb7832008-01-11 16:17:00 +0000784 __slots__ = ()
Christian Heimes043d6f62008-01-07 17:19:16 +0000785 _make = classmethod(tuple.__new__)
786 def _replace(self, _map=map, **kwds):
Christian Heimes2380ac72008-01-09 00:17:24 +0000787 return self._make(_map(kwds.get, ('x', 'y'), self))
Christian Heimes043d6f62008-01-07 17:19:16 +0000788
789 print(Point(11, 22)._replace(x=100))
Guido van Rossum3d392eb2007-11-16 00:35:22 +0000790
Christian Heimes25bb7832008-01-11 16:17:00 +0000791 Point3D = namedtuple('Point3D', Point._fields + ('z',))
792 print(Point3D.__doc__)
793
Guido van Rossumd8faa362007-04-27 19:54:29 +0000794 import doctest
Guido van Rossum8ce8a782007-11-01 19:42:39 +0000795 TestResults = namedtuple('TestResults', 'failed attempted')
Guido van Rossumd8faa362007-04-27 19:54:29 +0000796 print(TestResults(*doctest.testmod()))