blob: eaecd967ba93f6d845f80c7434ac470d7700a340 [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 Hettingerf1736542009-03-23 05:19:21 +000014from weakref import 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 Hettingerf1736542009-03-23 05:19:21 +000044 self.__root = _Link() # sentinel node for the doubly linked list
Raymond Hettingerdc879f02009-03-19 20:30:56 +000045 self.clear()
Raymond Hettinger2d32f632009-03-02 21:24:57 +000046 self.update(*args, **kwds)
47
48 def clear(self):
Raymond Hettingerf1736542009-03-23 05:19:21 +000049 root = self.__root
50 root.prev = root.next = root
51 self.__map = {}
Raymond Hettinger2d32f632009-03-02 21:24:57 +000052 dict.clear(self)
53
54 def __setitem__(self, key, value):
Raymond Hettingerf1736542009-03-23 05:19:21 +000055 # Setting a new item creates a new link which goes at the end of the linked
56 # list, and the inherited dictionary is updated with the new key/value pair.
Raymond Hettinger2d32f632009-03-02 21:24:57 +000057 if key not in self:
Raymond Hettingerf1736542009-03-23 05:19:21 +000058 self.__map[key] = link = _Link()
59 root = self.__root
60 last = root.prev
61 link.prev, link.next, link.key = last, root, key
62 last.next = root.prev = proxy(link)
Raymond Hettinger2d32f632009-03-02 21:24:57 +000063 dict.__setitem__(self, key, value)
64
65 def __delitem__(self, key):
Raymond Hettingerf1736542009-03-23 05:19:21 +000066 # Deleting an existing item uses self.__map to find the link which is
67 # then removed by updating the links in the predecessor and successor nodes.
Raymond Hettinger2d32f632009-03-02 21:24:57 +000068 dict.__delitem__(self, key)
Raymond Hettingerf1736542009-03-23 05:19:21 +000069 link = self.__map.pop(key)
70 link.prev.next = link.next
71 link.next.prev = link.prev
Raymond Hettinger2d32f632009-03-02 21:24:57 +000072
73 def __iter__(self):
Raymond Hettingerf1736542009-03-23 05:19:21 +000074 # Traverse the linked list in order.
75 root = self.__root
76 curr = root.next
77 while curr is not root:
78 yield curr.key
79 curr = curr.next
Raymond Hettinger2d32f632009-03-02 21:24:57 +000080
81 def __reversed__(self):
Raymond Hettingerf1736542009-03-23 05:19:21 +000082 # Traverse the linked list in reverse order.
83 root = self.__root
84 curr = root.prev
85 while curr is not root:
86 yield curr.key
87 curr = curr.prev
Raymond Hettinger2d32f632009-03-02 21:24:57 +000088
89 def __reduce__(self):
90 items = [[k, self[k]] for k in self]
Raymond Hettingerf1736542009-03-23 05:19:21 +000091 tmp = self.__map, self.__root
92 del self.__map, self.__root
Raymond Hettinger2d32f632009-03-02 21:24:57 +000093 inst_dict = vars(self).copy()
Raymond Hettingerf1736542009-03-23 05:19:21 +000094 self.__map, self.__root = tmp
Raymond Hettinger14b89ff2009-03-03 22:20:56 +000095 if inst_dict:
96 return (self.__class__, (items,), inst_dict)
97 return self.__class__, (items,)
Raymond Hettinger2d32f632009-03-02 21:24:57 +000098
99 setdefault = MutableMapping.setdefault
100 update = MutableMapping.update
101 pop = MutableMapping.pop
102 keys = MutableMapping.keys
103 values = MutableMapping.values
104 items = MutableMapping.items
105
Raymond Hettinger18ed2cb2009-03-19 23:14:39 +0000106 def popitem(self, last=True):
107 if not self:
108 raise KeyError('dictionary is empty')
109 key = next(reversed(self)) if last else next(iter(self))
110 value = self.pop(key)
111 return key, value
112
Raymond Hettinger2d32f632009-03-02 21:24:57 +0000113 def __repr__(self):
114 if not self:
115 return '%s()' % (self.__class__.__name__,)
116 return '%s(%r)' % (self.__class__.__name__, list(self.items()))
117
118 def copy(self):
119 return self.__class__(self)
120
121 @classmethod
122 def fromkeys(cls, iterable, value=None):
123 d = cls()
124 for key in iterable:
125 d[key] = value
126 return d
127
128 def __eq__(self, other):
129 if isinstance(other, OrderedDict):
Raymond Hettingerea9f8db2009-03-02 21:28:41 +0000130 return len(self)==len(other) and all(p==q for p, q in zip(self.items(), other.items()))
Raymond Hettinger2d32f632009-03-02 21:24:57 +0000131 return dict.__eq__(self, other)
132
Raymond Hettingerb8baf632009-01-14 02:20:07 +0000133
Christian Heimes99170a52007-12-19 02:07:34 +0000134
Raymond Hettinger48b8b662008-02-05 01:53:00 +0000135################################################################################
136### namedtuple
137################################################################################
138
Benjamin Petersona86f2c02009-02-10 02:41:10 +0000139def namedtuple(typename, field_names, verbose=False, rename=False):
Guido van Rossumd8faa362007-04-27 19:54:29 +0000140 """Returns a new subclass of tuple with named fields.
141
Guido van Rossum8ce8a782007-11-01 19:42:39 +0000142 >>> Point = namedtuple('Point', 'x y')
Thomas Wouters1b7f8912007-09-19 03:06:30 +0000143 >>> Point.__doc__ # docstring for the new class
Guido van Rossumd8faa362007-04-27 19:54:29 +0000144 'Point(x, y)'
Thomas Wouters1b7f8912007-09-19 03:06:30 +0000145 >>> p = Point(11, y=22) # instantiate with positional args or keywords
Christian Heimes99170a52007-12-19 02:07:34 +0000146 >>> p[0] + p[1] # indexable like a plain tuple
Guido van Rossumd8faa362007-04-27 19:54:29 +0000147 33
Christian Heimes99170a52007-12-19 02:07:34 +0000148 >>> x, y = p # unpack like a regular tuple
Guido van Rossumd8faa362007-04-27 19:54:29 +0000149 >>> x, y
150 (11, 22)
Thomas Wouters1b7f8912007-09-19 03:06:30 +0000151 >>> p.x + p.y # fields also accessable by name
Guido van Rossumd8faa362007-04-27 19:54:29 +0000152 33
Christian Heimes0449f632007-12-15 01:27:15 +0000153 >>> d = p._asdict() # convert to a dictionary
Guido van Rossum8ce8a782007-11-01 19:42:39 +0000154 >>> d['x']
155 11
156 >>> Point(**d) # convert from a dictionary
Guido van Rossumd8faa362007-04-27 19:54:29 +0000157 Point(x=11, y=22)
Christian Heimes0449f632007-12-15 01:27:15 +0000158 >>> p._replace(x=100) # _replace() is like str.replace() but targets named fields
Thomas Wouters1b7f8912007-09-19 03:06:30 +0000159 Point(x=100, y=22)
Guido van Rossumd8faa362007-04-27 19:54:29 +0000160
161 """
162
Christian Heimes2380ac72008-01-09 00:17:24 +0000163 # Parse and validate the field names. Validation serves two purposes,
164 # generating informative error messages and preventing template injection attacks.
Guido van Rossum8ce8a782007-11-01 19:42:39 +0000165 if isinstance(field_names, str):
166 field_names = field_names.replace(',', ' ').split() # names separated by whitespace and/or commas
Benjamin Petersone9bbc8b2008-09-28 02:06:32 +0000167 field_names = tuple(map(str, field_names))
Benjamin Petersona86f2c02009-02-10 02:41:10 +0000168 if rename:
169 names = list(field_names)
170 seen = set()
171 for i, name in enumerate(names):
172 if (not all(c.isalnum() or c=='_' for c in name) or _iskeyword(name)
173 or not name or name[0].isdigit() or name.startswith('_')
174 or name in seen):
175 names[i] = '_%d' % (i+1)
176 seen.add(name)
177 field_names = tuple(names)
Guido van Rossum8ce8a782007-11-01 19:42:39 +0000178 for name in (typename,) + field_names:
Christian Heimesb9eccbf2007-12-05 20:18:38 +0000179 if not all(c.isalnum() or c=='_' for c in name):
Guido van Rossum8ce8a782007-11-01 19:42:39 +0000180 raise ValueError('Type names and field names can only contain alphanumeric characters and underscores: %r' % name)
181 if _iskeyword(name):
182 raise ValueError('Type names and field names cannot be a keyword: %r' % name)
183 if name[0].isdigit():
184 raise ValueError('Type names and field names cannot start with a number: %r' % name)
185 seen_names = set()
186 for name in field_names:
Benjamin Petersona86f2c02009-02-10 02:41:10 +0000187 if name.startswith('_') and not rename:
Christian Heimes0449f632007-12-15 01:27:15 +0000188 raise ValueError('Field names cannot start with an underscore: %r' % name)
Guido van Rossum8ce8a782007-11-01 19:42:39 +0000189 if name in seen_names:
190 raise ValueError('Encountered duplicate field name: %r' % name)
191 seen_names.add(name)
192
193 # Create and fill-in the class template
Christian Heimesfaf2f632008-01-06 16:59:19 +0000194 numfields = len(field_names)
Thomas Wouters1b7f8912007-09-19 03:06:30 +0000195 argtxt = repr(field_names).replace("'", "")[1:-1] # tuple repr without parens or quotes
Guido van Rossumd59da4b2007-05-22 18:11:13 +0000196 reprtxt = ', '.join('%s=%%r' % name for name in field_names)
197 template = '''class %(typename)s(tuple):
Christian Heimes0449f632007-12-15 01:27:15 +0000198 '%(typename)s(%(argtxt)s)' \n
199 __slots__ = () \n
Christian Heimesfaf2f632008-01-06 16:59:19 +0000200 _fields = %(field_names)r \n
Guido van Rossumd59da4b2007-05-22 18:11:13 +0000201 def __new__(cls, %(argtxt)s):
Christian Heimes0449f632007-12-15 01:27:15 +0000202 return tuple.__new__(cls, (%(argtxt)s)) \n
Christian Heimesfaf2f632008-01-06 16:59:19 +0000203 @classmethod
Christian Heimes043d6f62008-01-07 17:19:16 +0000204 def _make(cls, iterable, new=tuple.__new__, len=len):
Christian Heimesfaf2f632008-01-06 16:59:19 +0000205 'Make a new %(typename)s object from a sequence or iterable'
Christian Heimes043d6f62008-01-07 17:19:16 +0000206 result = new(cls, iterable)
Christian Heimesfaf2f632008-01-06 16:59:19 +0000207 if len(result) != %(numfields)d:
208 raise TypeError('Expected %(numfields)d arguments, got %%d' %% len(result))
209 return result \n
Guido van Rossumd59da4b2007-05-22 18:11:13 +0000210 def __repr__(self):
Christian Heimes0449f632007-12-15 01:27:15 +0000211 return '%(typename)s(%(reprtxt)s)' %% self \n
Raymond Hettingera4f52b12009-03-02 22:28:31 +0000212 def _asdict(self):
213 'Return a new OrderedDict which maps field names to their values'
214 return OrderedDict(zip(self._fields, self)) \n
Christian Heimes0449f632007-12-15 01:27:15 +0000215 def _replace(self, **kwds):
Guido van Rossum3d392eb2007-11-16 00:35:22 +0000216 'Return a new %(typename)s object replacing specified fields with new values'
Christian Heimesfaf2f632008-01-06 16:59:19 +0000217 result = self._make(map(kwds.pop, %(field_names)r, self))
218 if kwds:
219 raise ValueError('Got unexpected field names: %%r' %% kwds.keys())
Georg Brandlc28e1fa2008-06-10 19:20:26 +0000220 return result \n
221 def __getnewargs__(self):
222 return tuple(self) \n\n''' % locals()
Guido van Rossumd59da4b2007-05-22 18:11:13 +0000223 for i, name in enumerate(field_names):
Thomas Wouters1b7f8912007-09-19 03:06:30 +0000224 template += ' %s = property(itemgetter(%d))\n' % (name, i)
225 if verbose:
226 print(template)
Guido van Rossum8ce8a782007-11-01 19:42:39 +0000227
Georg Brandlf08a9dd2008-06-10 16:57:31 +0000228 # Execute the template string in a temporary namespace and
229 # support tracing utilities by setting a value for frame.f_globals['__name__']
Raymond Hettingera4f52b12009-03-02 22:28:31 +0000230 namespace = dict(itemgetter=_itemgetter, __name__='namedtuple_%s' % typename,
231 OrderedDict=OrderedDict)
Guido van Rossum8ce8a782007-11-01 19:42:39 +0000232 try:
233 exec(template, namespace)
234 except SyntaxError as e:
Christian Heimes99170a52007-12-19 02:07:34 +0000235 raise SyntaxError(e.msg + ':\n' + template) from e
Guido van Rossum8ce8a782007-11-01 19:42:39 +0000236 result = namespace[typename]
237
238 # For pickling to work, the __module__ variable needs to be set to the frame
239 # where the named tuple is created. Bypass this step in enviroments where
240 # sys._getframe is not defined (Jython for example).
Guido van Rossumd59da4b2007-05-22 18:11:13 +0000241 if hasattr(_sys, '_getframe'):
Raymond Hettinger0f055172009-01-27 10:06:09 +0000242 result.__module__ = _sys._getframe(1).f_globals.get('__name__', '__main__')
Guido van Rossum8ce8a782007-11-01 19:42:39 +0000243
Guido van Rossumd59da4b2007-05-22 18:11:13 +0000244 return result
Guido van Rossumd8faa362007-04-27 19:54:29 +0000245
Guido van Rossumd8faa362007-04-27 19:54:29 +0000246
Raymond Hettingerb8baf632009-01-14 02:20:07 +0000247########################################################################
248### Counter
249########################################################################
250
251class Counter(dict):
252 '''Dict subclass for counting hashable items. Sometimes called a bag
253 or multiset. Elements are stored as dictionary keys and their counts
254 are stored as dictionary values.
255
256 >>> c = Counter('abracadabra') # count elements from a string
257
258 >>> c.most_common(3) # three most common elements
259 [('a', 5), ('r', 2), ('b', 2)]
260 >>> sorted(c) # list all unique elements
261 ['a', 'b', 'c', 'd', 'r']
262 >>> ''.join(sorted(c.elements())) # list elements with repetitions
263 'aaaaabbcdrr'
264 >>> sum(c.values()) # total of all counts
265 11
266
267 >>> c['a'] # count of letter 'a'
268 5
269 >>> for elem in 'shazam': # update counts from an iterable
270 ... c[elem] += 1 # by adding 1 to each element's count
271 >>> c['a'] # now there are seven 'a'
272 7
273 >>> del c['r'] # remove all 'r'
274 >>> c['r'] # now there are zero 'r'
275 0
276
277 >>> d = Counter('simsalabim') # make another counter
278 >>> c.update(d) # add in the second counter
279 >>> c['a'] # now there are nine 'a'
280 9
281
282 >>> c.clear() # empty the counter
283 >>> c
284 Counter()
285
286 Note: If a count is set to zero or reduced to zero, it will remain
287 in the counter until the entry is deleted or the counter is cleared:
288
289 >>> c = Counter('aaabbc')
290 >>> c['b'] -= 2 # reduce the count of 'b' by two
291 >>> c.most_common() # 'b' is still in, but its count is zero
292 [('a', 3), ('c', 1), ('b', 0)]
293
294 '''
295 # References:
296 # http://en.wikipedia.org/wiki/Multiset
297 # http://www.gnu.org/software/smalltalk/manual-base/html_node/Bag.html
298 # http://www.demo2s.com/Tutorial/Cpp/0380__set-multiset/Catalog0380__set-multiset.htm
299 # http://code.activestate.com/recipes/259174/
300 # Knuth, TAOCP Vol. II section 4.6.3
301
Raymond Hettinger4d2073a2009-01-20 03:41:22 +0000302 def __init__(self, iterable=None, **kwds):
Raymond Hettingerb8baf632009-01-14 02:20:07 +0000303 '''Create a new, empty Counter object. And if given, count elements
304 from an input iterable. Or, initialize the count from another mapping
305 of elements to their counts.
306
307 >>> c = Counter() # a new, empty counter
308 >>> c = Counter('gallahad') # a new counter from an iterable
309 >>> c = Counter({'a': 4, 'b': 2}) # a new counter from a mapping
Raymond Hettinger4d2073a2009-01-20 03:41:22 +0000310 >>> c = Counter(a=4, b=2) # a new counter from keyword args
Raymond Hettingerb8baf632009-01-14 02:20:07 +0000311
312 '''
Raymond Hettinger4d2073a2009-01-20 03:41:22 +0000313 self.update(iterable, **kwds)
Raymond Hettingerb8baf632009-01-14 02:20:07 +0000314
315 def __missing__(self, key):
316 'The count of elements not in the Counter is zero.'
317 # Needed so that self[missing_item] does not raise KeyError
318 return 0
319
320 def most_common(self, n=None):
321 '''List the n most common elements and their counts from the most
322 common to the least. If n is None, then list all element counts.
323
324 >>> Counter('abracadabra').most_common(3)
325 [('a', 5), ('r', 2), ('b', 2)]
326
327 '''
328 # Emulate Bag.sortedByCount from Smalltalk
329 if n is None:
330 return sorted(self.items(), key=_itemgetter(1), reverse=True)
331 return _heapq.nlargest(n, self.items(), key=_itemgetter(1))
332
333 def elements(self):
334 '''Iterator over elements repeating each as many times as its count.
335
336 >>> c = Counter('ABCABC')
337 >>> sorted(c.elements())
338 ['A', 'A', 'B', 'B', 'C', 'C']
339
340 # Knuth's example for prime factors of 1836: 2**2 * 3**3 * 17**1
341 >>> prime_factors = Counter({2: 2, 3: 3, 17: 1})
342 >>> product = 1
343 >>> for factor in prime_factors.elements(): # loop over factors
344 ... product *= factor # and multiply them
345 >>> product
346 1836
347
348 Note, if an element's count has been set to zero or is a negative
349 number, elements() will ignore it.
350
351 '''
352 # Emulate Bag.do from Smalltalk and Multiset.begin from C++.
353 return _chain.from_iterable(_starmap(_repeat, self.items()))
354
355 # Override dict methods where necessary
356
357 @classmethod
358 def fromkeys(cls, iterable, v=None):
359 # There is no equivalent method for counters because setting v=1
360 # means that no element can have a count greater than one.
361 raise NotImplementedError(
362 'Counter.fromkeys() is undefined. Use Counter(iterable) instead.')
363
Raymond Hettinger4d2073a2009-01-20 03:41:22 +0000364 def update(self, iterable=None, **kwds):
Raymond Hettingerb8baf632009-01-14 02:20:07 +0000365 '''Like dict.update() but add counts instead of replacing them.
366
367 Source can be an iterable, a dictionary, or another Counter instance.
368
369 >>> c = Counter('which')
370 >>> c.update('witch') # add elements from another iterable
371 >>> d = Counter('watch')
372 >>> c.update(d) # add elements from another counter
373 >>> c['h'] # four 'h' in which, witch, and watch
374 4
375
376 '''
377 # The regular dict.update() operation makes no sense here because the
378 # replace behavior results in the some of original untouched counts
379 # being mixed-in with all of the other counts for a mismash that
380 # doesn't have a straight-forward interpretation in most counting
Raymond Hettinger4d2073a2009-01-20 03:41:22 +0000381 # contexts. Instead, we implement straight-addition. Both the inputs
382 # and outputs are allowed to contain zero and negative counts.
Raymond Hettingerb8baf632009-01-14 02:20:07 +0000383
384 if iterable is not None:
385 if isinstance(iterable, Mapping):
Raymond Hettingerdd01f8f2009-01-22 09:09:55 +0000386 if self:
387 for elem, count in iterable.items():
388 self[elem] += count
389 else:
390 dict.update(self, iterable) # fast path when counter is empty
Raymond Hettingerb8baf632009-01-14 02:20:07 +0000391 else:
392 for elem in iterable:
393 self[elem] += 1
Raymond Hettinger4d2073a2009-01-20 03:41:22 +0000394 if kwds:
395 self.update(kwds)
Raymond Hettingerb8baf632009-01-14 02:20:07 +0000396
397 def copy(self):
398 'Like dict.copy() but returns a Counter instance instead of a dict.'
399 return Counter(self)
400
Raymond Hettinger4d2073a2009-01-20 03:41:22 +0000401 def __delitem__(self, elem):
402 'Like dict.__delitem__() but does not raise KeyError for missing values.'
403 if elem in self:
404 dict.__delitem__(self, elem)
405
Raymond Hettingerb8baf632009-01-14 02:20:07 +0000406 def __repr__(self):
407 if not self:
408 return '%s()' % self.__class__.__name__
409 items = ', '.join(map('%r: %r'.__mod__, self.most_common()))
410 return '%s({%s})' % (self.__class__.__name__, items)
411
Raymond Hettinger4d2073a2009-01-20 03:41:22 +0000412 # Multiset-style mathematical operations discussed in:
413 # Knuth TAOCP Volume II section 4.6.3 exercise 19
414 # and at http://en.wikipedia.org/wiki/Multiset
415 #
Raymond Hettinger4d2073a2009-01-20 03:41:22 +0000416 # Outputs guaranteed to only include positive counts.
417 #
418 # To strip negative and zero counts, add-in an empty counter:
419 # c += Counter()
420
421 def __add__(self, other):
422 '''Add counts from two counters.
423
424 >>> Counter('abbb') + Counter('bcc')
425 Counter({'b': 4, 'c': 2, 'a': 1})
426
427 '''
428 if not isinstance(other, Counter):
429 return NotImplemented
430 result = Counter()
431 for elem in set(self) | set(other):
432 newcount = self[elem] + other[elem]
433 if newcount > 0:
434 result[elem] = newcount
435 return result
436
437 def __sub__(self, other):
438 ''' Subtract count, but keep only results with positive counts.
439
440 >>> Counter('abbbc') - Counter('bccd')
441 Counter({'b': 2, 'a': 1})
442
443 '''
444 if not isinstance(other, Counter):
445 return NotImplemented
446 result = Counter()
Raymond Hettingere0d1b9f2009-01-21 20:36:27 +0000447 for elem in set(self) | set(other):
448 newcount = self[elem] - other[elem]
Raymond Hettinger4d2073a2009-01-20 03:41:22 +0000449 if newcount > 0:
450 result[elem] = newcount
451 return result
452
453 def __or__(self, other):
454 '''Union is the maximum of value in either of the input counters.
455
456 >>> Counter('abbb') | Counter('bcc')
457 Counter({'b': 3, 'c': 2, 'a': 1})
458
459 '''
460 if not isinstance(other, Counter):
461 return NotImplemented
462 _max = max
463 result = Counter()
464 for elem in set(self) | set(other):
465 newcount = _max(self[elem], other[elem])
466 if newcount > 0:
467 result[elem] = newcount
468 return result
469
470 def __and__(self, other):
471 ''' Intersection is the minimum of corresponding counts.
472
473 >>> Counter('abbb') & Counter('bcc')
474 Counter({'b': 1})
475
476 '''
477 if not isinstance(other, Counter):
478 return NotImplemented
479 _min = min
480 result = Counter()
481 if len(self) < len(other):
482 self, other = other, self
483 for elem in filter(self.__contains__, other):
484 newcount = _min(self[elem], other[elem])
485 if newcount > 0:
486 result[elem] = newcount
487 return result
488
Guido van Rossumd8faa362007-04-27 19:54:29 +0000489
Raymond Hettinger48b8b662008-02-05 01:53:00 +0000490################################################################################
491### UserDict
492################################################################################
Guido van Rossumd8faa362007-04-27 19:54:29 +0000493
Raymond Hettinger48b8b662008-02-05 01:53:00 +0000494class UserDict(MutableMapping):
495
496 # Start by filling-out the abstract methods
497 def __init__(self, dict=None, **kwargs):
498 self.data = {}
499 if dict is not None:
500 self.update(dict)
501 if len(kwargs):
502 self.update(kwargs)
503 def __len__(self): return len(self.data)
504 def __getitem__(self, key):
505 if key in self.data:
506 return self.data[key]
507 if hasattr(self.__class__, "__missing__"):
508 return self.__class__.__missing__(self, key)
509 raise KeyError(key)
510 def __setitem__(self, key, item): self.data[key] = item
511 def __delitem__(self, key): del self.data[key]
512 def __iter__(self):
513 return iter(self.data)
514
Raymond Hettinger554c8b82008-02-05 22:54:43 +0000515 # Modify __contains__ to work correctly when __missing__ is present
516 def __contains__(self, key):
517 return key in self.data
Raymond Hettinger48b8b662008-02-05 01:53:00 +0000518
519 # Now, add the methods in dicts but not in MutableMapping
520 def __repr__(self): return repr(self.data)
521 def copy(self):
522 if self.__class__ is UserDict:
523 return UserDict(self.data.copy())
524 import copy
525 data = self.data
526 try:
527 self.data = {}
528 c = copy.copy(self)
529 finally:
530 self.data = data
531 c.update(self)
532 return c
533 @classmethod
534 def fromkeys(cls, iterable, value=None):
535 d = cls()
536 for key in iterable:
537 d[key] = value
538 return d
539
Raymond Hettinger48b8b662008-02-05 01:53:00 +0000540
541
542################################################################################
Raymond Hettinger53dbe392008-02-12 20:03:09 +0000543### UserList
544################################################################################
545
546class UserList(MutableSequence):
547 """A more or less complete user-defined wrapper around list objects."""
548 def __init__(self, initlist=None):
549 self.data = []
550 if initlist is not None:
551 # XXX should this accept an arbitrary sequence?
552 if type(initlist) == type(self.data):
553 self.data[:] = initlist
554 elif isinstance(initlist, UserList):
555 self.data[:] = initlist.data[:]
556 else:
557 self.data = list(initlist)
558 def __repr__(self): return repr(self.data)
559 def __lt__(self, other): return self.data < self.__cast(other)
560 def __le__(self, other): return self.data <= self.__cast(other)
561 def __eq__(self, other): return self.data == self.__cast(other)
562 def __ne__(self, other): return self.data != self.__cast(other)
563 def __gt__(self, other): return self.data > self.__cast(other)
564 def __ge__(self, other): return self.data >= self.__cast(other)
565 def __cast(self, other):
566 return other.data if isinstance(other, UserList) else other
Raymond Hettinger53dbe392008-02-12 20:03:09 +0000567 def __contains__(self, item): return item in self.data
568 def __len__(self): return len(self.data)
569 def __getitem__(self, i): return self.data[i]
570 def __setitem__(self, i, item): self.data[i] = item
571 def __delitem__(self, i): del self.data[i]
572 def __add__(self, other):
573 if isinstance(other, UserList):
574 return self.__class__(self.data + other.data)
575 elif isinstance(other, type(self.data)):
576 return self.__class__(self.data + other)
577 return self.__class__(self.data + list(other))
578 def __radd__(self, other):
579 if isinstance(other, UserList):
580 return self.__class__(other.data + self.data)
581 elif isinstance(other, type(self.data)):
582 return self.__class__(other + self.data)
583 return self.__class__(list(other) + self.data)
584 def __iadd__(self, other):
585 if isinstance(other, UserList):
586 self.data += other.data
587 elif isinstance(other, type(self.data)):
588 self.data += other
589 else:
590 self.data += list(other)
591 return self
592 def __mul__(self, n):
593 return self.__class__(self.data*n)
594 __rmul__ = __mul__
595 def __imul__(self, n):
596 self.data *= n
597 return self
598 def append(self, item): self.data.append(item)
599 def insert(self, i, item): self.data.insert(i, item)
600 def pop(self, i=-1): return self.data.pop(i)
601 def remove(self, item): self.data.remove(item)
602 def count(self, item): return self.data.count(item)
603 def index(self, item, *args): return self.data.index(item, *args)
604 def reverse(self): self.data.reverse()
605 def sort(self, *args, **kwds): self.data.sort(*args, **kwds)
606 def extend(self, other):
607 if isinstance(other, UserList):
608 self.data.extend(other.data)
609 else:
610 self.data.extend(other)
611
612
613
614################################################################################
Raymond Hettingerb3a65f82008-02-21 22:11:37 +0000615### UserString
616################################################################################
617
618class UserString(Sequence):
619 def __init__(self, seq):
620 if isinstance(seq, str):
621 self.data = seq
622 elif isinstance(seq, UserString):
623 self.data = seq.data[:]
624 else:
625 self.data = str(seq)
626 def __str__(self): return str(self.data)
627 def __repr__(self): return repr(self.data)
628 def __int__(self): return int(self.data)
Raymond Hettingerb3a65f82008-02-21 22:11:37 +0000629 def __float__(self): return float(self.data)
630 def __complex__(self): return complex(self.data)
631 def __hash__(self): return hash(self.data)
632
633 def __eq__(self, string):
634 if isinstance(string, UserString):
635 return self.data == string.data
636 return self.data == string
637 def __ne__(self, string):
638 if isinstance(string, UserString):
639 return self.data != string.data
640 return self.data != string
641 def __lt__(self, string):
642 if isinstance(string, UserString):
643 return self.data < string.data
644 return self.data < string
645 def __le__(self, string):
646 if isinstance(string, UserString):
647 return self.data <= string.data
648 return self.data <= string
649 def __gt__(self, string):
650 if isinstance(string, UserString):
651 return self.data > string.data
652 return self.data > string
653 def __ge__(self, string):
654 if isinstance(string, UserString):
655 return self.data >= string.data
656 return self.data >= string
657
658 def __contains__(self, char):
659 if isinstance(char, UserString):
660 char = char.data
661 return char in self.data
662
663 def __len__(self): return len(self.data)
664 def __getitem__(self, index): return self.__class__(self.data[index])
665 def __add__(self, other):
666 if isinstance(other, UserString):
667 return self.__class__(self.data + other.data)
668 elif isinstance(other, str):
669 return self.__class__(self.data + other)
670 return self.__class__(self.data + str(other))
671 def __radd__(self, other):
672 if isinstance(other, str):
673 return self.__class__(other + self.data)
674 return self.__class__(str(other) + self.data)
675 def __mul__(self, n):
676 return self.__class__(self.data*n)
677 __rmul__ = __mul__
678 def __mod__(self, args):
679 return self.__class__(self.data % args)
680
681 # the following methods are defined in alphabetical order:
682 def capitalize(self): return self.__class__(self.data.capitalize())
683 def center(self, width, *args):
684 return self.__class__(self.data.center(width, *args))
685 def count(self, sub, start=0, end=_sys.maxsize):
686 if isinstance(sub, UserString):
687 sub = sub.data
688 return self.data.count(sub, start, end)
689 def encode(self, encoding=None, errors=None): # XXX improve this?
690 if encoding:
691 if errors:
692 return self.__class__(self.data.encode(encoding, errors))
693 return self.__class__(self.data.encode(encoding))
694 return self.__class__(self.data.encode())
695 def endswith(self, suffix, start=0, end=_sys.maxsize):
696 return self.data.endswith(suffix, start, end)
697 def expandtabs(self, tabsize=8):
698 return self.__class__(self.data.expandtabs(tabsize))
699 def find(self, sub, start=0, end=_sys.maxsize):
700 if isinstance(sub, UserString):
701 sub = sub.data
702 return self.data.find(sub, start, end)
703 def format(self, *args, **kwds):
704 return self.data.format(*args, **kwds)
705 def index(self, sub, start=0, end=_sys.maxsize):
706 return self.data.index(sub, start, end)
707 def isalpha(self): return self.data.isalpha()
708 def isalnum(self): return self.data.isalnum()
709 def isdecimal(self): return self.data.isdecimal()
710 def isdigit(self): return self.data.isdigit()
711 def isidentifier(self): return self.data.isidentifier()
712 def islower(self): return self.data.islower()
713 def isnumeric(self): return self.data.isnumeric()
714 def isspace(self): return self.data.isspace()
715 def istitle(self): return self.data.istitle()
716 def isupper(self): return self.data.isupper()
717 def join(self, seq): return self.data.join(seq)
718 def ljust(self, width, *args):
719 return self.__class__(self.data.ljust(width, *args))
720 def lower(self): return self.__class__(self.data.lower())
721 def lstrip(self, chars=None): return self.__class__(self.data.lstrip(chars))
722 def partition(self, sep):
723 return self.data.partition(sep)
724 def replace(self, old, new, maxsplit=-1):
725 if isinstance(old, UserString):
726 old = old.data
727 if isinstance(new, UserString):
728 new = new.data
729 return self.__class__(self.data.replace(old, new, maxsplit))
730 def rfind(self, sub, start=0, end=_sys.maxsize):
731 return self.data.rfind(sub, start, end)
732 def rindex(self, sub, start=0, end=_sys.maxsize):
733 return self.data.rindex(sub, start, end)
734 def rjust(self, width, *args):
735 return self.__class__(self.data.rjust(width, *args))
736 def rpartition(self, sep):
737 return self.data.rpartition(sep)
738 def rstrip(self, chars=None):
739 return self.__class__(self.data.rstrip(chars))
740 def split(self, sep=None, maxsplit=-1):
741 return self.data.split(sep, maxsplit)
742 def rsplit(self, sep=None, maxsplit=-1):
743 return self.data.rsplit(sep, maxsplit)
744 def splitlines(self, keepends=0): return self.data.splitlines(keepends)
745 def startswith(self, prefix, start=0, end=_sys.maxsize):
746 return self.data.startswith(prefix, start, end)
747 def strip(self, chars=None): return self.__class__(self.data.strip(chars))
748 def swapcase(self): return self.__class__(self.data.swapcase())
749 def title(self): return self.__class__(self.data.title())
750 def translate(self, *args):
751 return self.__class__(self.data.translate(*args))
752 def upper(self): return self.__class__(self.data.upper())
753 def zfill(self, width): return self.__class__(self.data.zfill(width))
754
755
756
757################################################################################
Raymond Hettinger48b8b662008-02-05 01:53:00 +0000758### Simple tests
759################################################################################
Guido van Rossumd8faa362007-04-27 19:54:29 +0000760
761if __name__ == '__main__':
Thomas Wouters1b7f8912007-09-19 03:06:30 +0000762 # verify that instances can be pickled
Guido van Rossum99603b02007-07-20 00:22:32 +0000763 from pickle import loads, dumps
Guido van Rossum8ce8a782007-11-01 19:42:39 +0000764 Point = namedtuple('Point', 'x, y', True)
Guido van Rossumd8faa362007-04-27 19:54:29 +0000765 p = Point(x=10, y=20)
766 assert p == loads(dumps(p))
767
Guido van Rossum3d392eb2007-11-16 00:35:22 +0000768 # test and demonstrate ability to override methods
Christian Heimes043d6f62008-01-07 17:19:16 +0000769 class Point(namedtuple('Point', 'x y')):
Christian Heimes25bb7832008-01-11 16:17:00 +0000770 __slots__ = ()
Christian Heimes043d6f62008-01-07 17:19:16 +0000771 @property
772 def hypot(self):
773 return (self.x ** 2 + self.y ** 2) ** 0.5
Christian Heimes790c8232008-01-07 21:14:23 +0000774 def __str__(self):
Christian Heimes25bb7832008-01-11 16:17:00 +0000775 return 'Point: x=%6.3f y=%6.3f hypot=%6.3f' % (self.x, self.y, self.hypot)
Christian Heimes043d6f62008-01-07 17:19:16 +0000776
Christian Heimes25bb7832008-01-11 16:17:00 +0000777 for p in Point(3, 4), Point(14, 5/7.):
Christian Heimes790c8232008-01-07 21:14:23 +0000778 print (p)
Christian Heimes043d6f62008-01-07 17:19:16 +0000779
780 class Point(namedtuple('Point', 'x y')):
781 'Point class with optimized _make() and _replace() without error-checking'
Christian Heimes25bb7832008-01-11 16:17:00 +0000782 __slots__ = ()
Christian Heimes043d6f62008-01-07 17:19:16 +0000783 _make = classmethod(tuple.__new__)
784 def _replace(self, _map=map, **kwds):
Christian Heimes2380ac72008-01-09 00:17:24 +0000785 return self._make(_map(kwds.get, ('x', 'y'), self))
Christian Heimes043d6f62008-01-07 17:19:16 +0000786
787 print(Point(11, 22)._replace(x=100))
Guido van Rossum3d392eb2007-11-16 00:35:22 +0000788
Christian Heimes25bb7832008-01-11 16:17:00 +0000789 Point3D = namedtuple('Point3D', Point._fields + ('z',))
790 print(Point3D.__doc__)
791
Guido van Rossumd8faa362007-04-27 19:54:29 +0000792 import doctest
Guido van Rossum8ce8a782007-11-01 19:42:39 +0000793 TestResults = namedtuple('TestResults', 'failed attempted')
Guido van Rossumd8faa362007-04-27 19:54:29 +0000794 print(TestResults(*doctest.testmod()))