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