blob: dfbd02370586302324fcc3a1d1beb4f4ea3d7280 [file] [log] [blame]
Thomas Wouters73e5a5b2006-06-08 15:35:45 +00001"""functools.py - Tools for working with functions and callable objects
Thomas Wouters4d70c3d2006-06-08 14:42:34 +00002"""
3# Python module wrapper for _functools C module
4# to allow utilities written in Python to be added
5# to the functools module.
6# Written by Nick Coghlan <ncoghlan at gmail.com>
Georg Brandl2e7346a2010-07-31 18:09:23 +00007# and Raymond Hettinger <python at rcn.com>
8# Copyright (C) 2006-2010 Python Software Foundation.
Thomas Wouters73e5a5b2006-06-08 15:35:45 +00009# See C source code for _functools credits/copyright
Thomas Wouters4d70c3d2006-06-08 14:42:34 +000010
Georg Brandl2e7346a2010-07-31 18:09:23 +000011__all__ = ['update_wrapper', 'wraps', 'WRAPPER_ASSIGNMENTS', 'WRAPPER_UPDATES',
Benjamin Peterson1017ae52010-09-10 23:35:52 +000012 'total_ordering', 'cmp_to_key', 'lru_cache', 'reduce', 'partial']
Georg Brandl2e7346a2010-07-31 18:09:23 +000013
Guido van Rossum0919a1a2006-08-26 20:49:04 +000014from _functools import partial, reduce
Raymond Hettingerec0e9102012-03-16 01:16:31 -070015from collections import namedtuple
Raymond Hettingercbe88132010-08-14 22:22:10 +000016try:
Raymond Hettingerfd541172013-03-01 03:47:57 -080017 from _thread import RLock
Raymond Hettingercbe88132010-08-14 22:22:10 +000018except:
Raymond Hettinger409f6632013-03-01 23:20:13 -080019 class RLock:
Raymond Hettingerf96b2b02013-03-08 21:11:55 -070020 'Dummy reentrant lock for builds without threads'
Raymond Hettinger409f6632013-03-01 23:20:13 -080021 def __enter__(self): pass
22 def __exit__(self, exctype, excinst, exctb): pass
Thomas Wouters4d70c3d2006-06-08 14:42:34 +000023
Raymond Hettingerbc8e81d2012-03-17 00:24:09 -070024
25################################################################################
26### update_wrapper() and wraps() decorator
27################################################################################
28
Thomas Wouters73e5a5b2006-06-08 15:35:45 +000029# update_wrapper() and wraps() are tools to help write
30# wrapper functions that can handle naive introspection
31
Meador Ingeff7f64c2011-12-11 22:37:31 -060032WRAPPER_ASSIGNMENTS = ('__module__', '__name__', '__qualname__', '__doc__',
33 '__annotations__')
Thomas Wouters73e5a5b2006-06-08 15:35:45 +000034WRAPPER_UPDATES = ('__dict__',)
35def update_wrapper(wrapper,
36 wrapped,
37 assigned = WRAPPER_ASSIGNMENTS,
38 updated = WRAPPER_UPDATES):
39 """Update a wrapper function to look like the wrapped function
40
41 wrapper is the function to be updated
42 wrapped is the original function
43 assigned is a tuple naming the attributes assigned directly
44 from the wrapped function to the wrapper function (defaults to
45 functools.WRAPPER_ASSIGNMENTS)
Thomas Wouters89f507f2006-12-13 04:49:30 +000046 updated is a tuple naming the attributes of the wrapper that
Thomas Wouters73e5a5b2006-06-08 15:35:45 +000047 are updated with the corresponding attribute from the wrapped
48 function (defaults to functools.WRAPPER_UPDATES)
49 """
Nick Coghlan98876832010-08-17 06:17:18 +000050 wrapper.__wrapped__ = wrapped
Thomas Wouters73e5a5b2006-06-08 15:35:45 +000051 for attr in assigned:
Nick Coghlan98876832010-08-17 06:17:18 +000052 try:
53 value = getattr(wrapped, attr)
54 except AttributeError:
55 pass
56 else:
57 setattr(wrapper, attr, value)
Thomas Wouters73e5a5b2006-06-08 15:35:45 +000058 for attr in updated:
Thomas Wouters89f507f2006-12-13 04:49:30 +000059 getattr(wrapper, attr).update(getattr(wrapped, attr, {}))
Thomas Wouters73e5a5b2006-06-08 15:35:45 +000060 # Return the wrapper so this can be used as a decorator via partial()
61 return wrapper
62
63def wraps(wrapped,
64 assigned = WRAPPER_ASSIGNMENTS,
65 updated = WRAPPER_UPDATES):
66 """Decorator factory to apply update_wrapper() to a wrapper function
67
68 Returns a decorator that invokes update_wrapper() with the decorated
69 function as the wrapper argument and the arguments to wraps() as the
70 remaining arguments. Default arguments are as for update_wrapper().
71 This is a convenience function to simplify applying partial() to
72 update_wrapper().
73 """
74 return partial(update_wrapper, wrapped=wrapped,
75 assigned=assigned, updated=updated)
Raymond Hettingerc50846a2010-04-05 18:56:31 +000076
Raymond Hettingerbc8e81d2012-03-17 00:24:09 -070077
78################################################################################
79### total_ordering class decorator
80################################################################################
81
Raymond Hettingerc50846a2010-04-05 18:56:31 +000082def total_ordering(cls):
Georg Brandle5a26732010-05-19 21:06:36 +000083 """Class decorator that fills in missing ordering methods"""
Raymond Hettingerc50846a2010-04-05 18:56:31 +000084 convert = {
Raymond Hettinger23f9fc32011-01-08 07:01:56 +000085 '__lt__': [('__gt__', lambda self, other: not (self < other or self == other)),
86 ('__le__', lambda self, other: self < other or self == other),
Raymond Hettingerc50846a2010-04-05 18:56:31 +000087 ('__ge__', lambda self, other: not self < other)],
Raymond Hettinger23f9fc32011-01-08 07:01:56 +000088 '__le__': [('__ge__', lambda self, other: not self <= other or self == other),
89 ('__lt__', lambda self, other: self <= other and not self == other),
Raymond Hettingerc50846a2010-04-05 18:56:31 +000090 ('__gt__', lambda self, other: not self <= other)],
Raymond Hettinger23f9fc32011-01-08 07:01:56 +000091 '__gt__': [('__lt__', lambda self, other: not (self > other or self == other)),
92 ('__ge__', lambda self, other: self > other or self == other),
Raymond Hettingerc50846a2010-04-05 18:56:31 +000093 ('__le__', lambda self, other: not self > other)],
Raymond Hettinger23f9fc32011-01-08 07:01:56 +000094 '__ge__': [('__le__', lambda self, other: (not self >= other) or self == other),
95 ('__gt__', lambda self, other: self >= other and not self == other),
Raymond Hettingerc50846a2010-04-05 18:56:31 +000096 ('__lt__', lambda self, other: not self >= other)]
97 }
Raymond Hettinger3255c632010-09-16 00:31:21 +000098 # Find user-defined comparisons (not those inherited from object).
Raymond Hettinger1006bd42010-09-14 22:55:13 +000099 roots = [op for op in convert if getattr(cls, op, None) is not getattr(object, op, None)]
Raymond Hettinger56de7e22010-04-10 16:59:03 +0000100 if not roots:
101 raise ValueError('must define at least one ordering operation: < > <= >=')
102 root = max(roots) # prefer __lt__ to __le__ to __gt__ to __ge__
Raymond Hettingerc50846a2010-04-05 18:56:31 +0000103 for opname, opfunc in convert[root]:
104 if opname not in roots:
105 opfunc.__name__ = opname
106 opfunc.__doc__ = getattr(int, opname).__doc__
107 setattr(cls, opname, opfunc)
108 return cls
109
Raymond Hettingerbc8e81d2012-03-17 00:24:09 -0700110
111################################################################################
112### cmp_to_key() function converter
113################################################################################
114
Raymond Hettingerc50846a2010-04-05 18:56:31 +0000115def cmp_to_key(mycmp):
Georg Brandle5a26732010-05-19 21:06:36 +0000116 """Convert a cmp= function into a key= function"""
Raymond Hettingerc50846a2010-04-05 18:56:31 +0000117 class K(object):
Raymond Hettingera0d1d962011-03-21 17:50:28 -0700118 __slots__ = ['obj']
Raymond Hettinger7ab9e222011-04-05 02:33:54 -0700119 def __init__(self, obj):
Raymond Hettingerc50846a2010-04-05 18:56:31 +0000120 self.obj = obj
121 def __lt__(self, other):
122 return mycmp(self.obj, other.obj) < 0
123 def __gt__(self, other):
124 return mycmp(self.obj, other.obj) > 0
125 def __eq__(self, other):
126 return mycmp(self.obj, other.obj) == 0
127 def __le__(self, other):
128 return mycmp(self.obj, other.obj) <= 0
129 def __ge__(self, other):
130 return mycmp(self.obj, other.obj) >= 0
131 def __ne__(self, other):
132 return mycmp(self.obj, other.obj) != 0
Raymond Hettinger003be522011-05-03 11:01:32 -0700133 __hash__ = None
Raymond Hettingerc50846a2010-04-05 18:56:31 +0000134 return K
Georg Brandl2e7346a2010-07-31 18:09:23 +0000135
Raymond Hettinger7ab9e222011-04-05 02:33:54 -0700136try:
137 from _functools import cmp_to_key
138except ImportError:
139 pass
140
Raymond Hettingerbc8e81d2012-03-17 00:24:09 -0700141
142################################################################################
143### LRU Cache function decorator
144################################################################################
145
Raymond Hettingerdce583e2012-03-16 22:12:20 -0700146_CacheInfo = namedtuple("CacheInfo", ["hits", "misses", "maxsize", "currsize"])
Nick Coghlan234515a2010-11-30 06:19:46 +0000147
Raymond Hettinger0c9050c2012-06-04 00:21:14 -0700148class _HashedSeq(list):
Raymond Hettingerf96b2b02013-03-08 21:11:55 -0700149 """ This class guarantees that hash() will be called no more than once
150 per element. This is important because the lru_cache() will hash
151 the key multiple times on a cache miss.
152
153 """
154
Raymond Hettinger9acbb602012-04-30 22:32:16 -0700155 __slots__ = 'hashvalue'
156
Raymond Hettinger0c9050c2012-06-04 00:21:14 -0700157 def __init__(self, tup, hash=hash):
158 self[:] = tup
159 self.hashvalue = hash(tup)
Raymond Hettinger9acbb602012-04-30 22:32:16 -0700160
161 def __hash__(self):
162 return self.hashvalue
163
Raymond Hettinger0c9050c2012-06-04 00:21:14 -0700164def _make_key(args, kwds, typed,
165 kwd_mark = (object(),),
166 fasttypes = {int, str, frozenset, type(None)},
167 sorted=sorted, tuple=tuple, type=type, len=len):
Raymond Hettingerf96b2b02013-03-08 21:11:55 -0700168 """Make a cache key from optionally typed positional and keyword arguments
169
170 The key is constructed in a way that is flat as possible rather than
171 as a nested structure that would take more memory.
172
173 If there is only a single argument and its data type is known to cache
174 its hash value, then that argument is returned without a wrapper. This
175 saves space and improves lookup speed.
176
177 """
Raymond Hettinger0c9050c2012-06-04 00:21:14 -0700178 key = args
179 if kwds:
180 sorted_items = sorted(kwds.items())
181 key += kwd_mark
182 for item in sorted_items:
183 key += item
184 if typed:
185 key += tuple(type(v) for v in args)
186 if kwds:
187 key += tuple(type(v) for k, v in sorted_items)
188 elif len(key) == 1 and type(key[0]) in fasttypes:
189 return key[0]
190 return _HashedSeq(key)
191
Raymond Hettinger010ce322012-05-19 21:20:48 -0700192def lru_cache(maxsize=128, typed=False):
Benjamin Peterson1f594ad2010-08-08 13:17:07 +0000193 """Least-recently-used cache decorator.
Georg Brandl2e7346a2010-07-31 18:09:23 +0000194
Raymond Hettingerc79fb0e2010-12-01 03:45:41 +0000195 If *maxsize* is set to None, the LRU features are disabled and the cache
196 can grow without bound.
197
Raymond Hettingercd9fdfd2011-10-20 08:57:45 -0700198 If *typed* is True, arguments of different types will be cached separately.
199 For example, f(3.0) and f(3) will be treated as distinct calls with
200 distinct results.
201
Georg Brandl2e7346a2010-07-31 18:09:23 +0000202 Arguments to the cached function must be hashable.
Raymond Hettinger5fa40c02010-11-25 08:11:57 +0000203
Raymond Hettinger7f7a5a72012-03-30 21:50:40 -0700204 View the cache statistics named tuple (hits, misses, maxsize, currsize)
205 with f.cache_info(). Clear the cache and statistics with f.cache_clear().
Raymond Hettinger00f2f972010-12-01 00:47:56 +0000206 Access the underlying function with f.__wrapped__.
Raymond Hettinger5fa40c02010-11-25 08:11:57 +0000207
208 See: http://en.wikipedia.org/wiki/Cache_algorithms#Least_Recently_Used
Georg Brandl2e7346a2010-07-31 18:09:23 +0000209
Benjamin Peterson1f594ad2010-08-08 13:17:07 +0000210 """
Raymond Hettinger1ff50df2012-03-30 13:15:48 -0700211
Raymond Hettinger5fa40c02010-11-25 08:11:57 +0000212 # Users should only access the lru_cache through its public API:
Raymond Hettinger5e20bab2010-11-30 07:13:04 +0000213 # cache_info, cache_clear, and f.__wrapped__
Raymond Hettinger5fa40c02010-11-25 08:11:57 +0000214 # The internals of the lru_cache are encapsulated for thread safety and
215 # to allow the implementation to change (including a possible C version).
216
Raymond Hettinger9f0ab9f2012-04-29 14:55:27 -0700217 # Constants shared by all lru cache instances:
Raymond Hettingerb6b98c02012-04-29 18:09:02 -0700218 sentinel = object() # unique object used to signal cache misses
Raymond Hettinger0c9050c2012-06-04 00:21:14 -0700219 make_key = _make_key # build a key from the function arguments
Raymond Hettinger9f0ab9f2012-04-29 14:55:27 -0700220 PREV, NEXT, KEY, RESULT = 0, 1, 2, 3 # names for the link fields
221
Raymond Hettinger6e8c8172012-03-16 16:53:05 -0700222 def decorating_function(user_function):
Raymond Hettinger5fa40c02010-11-25 08:11:57 +0000223
Raymond Hettinger7f7a5a72012-03-30 21:50:40 -0700224 cache = {}
Raymond Hettingerb6b98c02012-04-29 18:09:02 -0700225 hits = misses = currsize = 0
Raymond Hettinger018b4fb2012-04-30 20:48:55 -0700226 full = False
Raymond Hettingerc6897852012-03-31 02:19:06 -0700227 cache_get = cache.get # bound method to lookup a key or return None
Raymond Hettingerfd541172013-03-01 03:47:57 -0800228 lock = RLock() # because linkedlist updates aren't threadsafe
Raymond Hettinger7f7a5a72012-03-30 21:50:40 -0700229 root = [] # root of the circular doubly linked list
230 root[:] = [root, root, None, None] # initialize by pointing to self
Raymond Hettinger6e8c8172012-03-16 16:53:05 -0700231
Raymond Hettinger7e0c5812012-03-17 15:10:24 -0700232 if maxsize == 0:
233
Raymond Hettinger7e0c5812012-03-17 15:10:24 -0700234 def wrapper(*args, **kwds):
Raymond Hettingerf96b2b02013-03-08 21:11:55 -0700235 # No caching -- just a statistics update after a successful call
Raymond Hettinger7e0c5812012-03-17 15:10:24 -0700236 nonlocal misses
Raymond Hettinger7dabfed2012-03-17 15:11:09 -0700237 result = user_function(*args, **kwds)
Raymond Hettinger7e0c5812012-03-17 15:10:24 -0700238 misses += 1
239 return result
240
241 elif maxsize is None:
Raymond Hettingerbc8e81d2012-03-17 00:24:09 -0700242
Raymond Hettingerc79fb0e2010-12-01 03:45:41 +0000243 def wrapper(*args, **kwds):
Raymond Hettingerf96b2b02013-03-08 21:11:55 -0700244 # Simple caching without ordering or size limit
Raymond Hettingerb6b98c02012-04-29 18:09:02 -0700245 nonlocal hits, misses, currsize
Raymond Hettinger9acbb602012-04-30 22:32:16 -0700246 key = make_key(args, kwds, typed)
Raymond Hettinger7f7a5a72012-03-30 21:50:40 -0700247 result = cache_get(key, sentinel)
248 if result is not sentinel:
Nick Coghlan234515a2010-11-30 06:19:46 +0000249 hits += 1
Raymond Hettinger4b779b32011-10-15 23:50:42 -0700250 return result
Raymond Hettinger4b779b32011-10-15 23:50:42 -0700251 result = user_function(*args, **kwds)
252 cache[key] = result
253 misses += 1
Raymond Hettingerb6b98c02012-04-29 18:09:02 -0700254 currsize += 1
Raymond Hettingerc79fb0e2010-12-01 03:45:41 +0000255 return result
Raymond Hettingerbc8e81d2012-03-17 00:24:09 -0700256
Raymond Hettingerc79fb0e2010-12-01 03:45:41 +0000257 else:
Raymond Hettingerbc8e81d2012-03-17 00:24:09 -0700258
Raymond Hettingerc79fb0e2010-12-01 03:45:41 +0000259 def wrapper(*args, **kwds):
Raymond Hettingerf96b2b02013-03-08 21:11:55 -0700260 # Size limited caching that tracks accesses by recency
Raymond Hettinger018b4fb2012-04-30 20:48:55 -0700261 nonlocal root, hits, misses, currsize, full
Raymond Hettinger9acbb602012-04-30 22:32:16 -0700262 key = make_key(args, kwds, typed)
Raymond Hettinger4b779b32011-10-15 23:50:42 -0700263 with lock:
Raymond Hettingerec0e9102012-03-16 01:16:31 -0700264 link = cache_get(key)
265 if link is not None:
Raymond Hettingerf96b2b02013-03-08 21:11:55 -0700266 # Move the link to the front of the circular queue
267 link_prev, link_next, _key, result = link
Raymond Hettingerec0e9102012-03-16 01:16:31 -0700268 link_prev[NEXT] = link_next
269 link_next[PREV] = link_prev
270 last = root[PREV]
271 last[NEXT] = root[PREV] = link
272 link[PREV] = last
273 link[NEXT] = root
Raymond Hettingerc79fb0e2010-12-01 03:45:41 +0000274 hits += 1
Raymond Hettinger4b779b32011-10-15 23:50:42 -0700275 return result
Raymond Hettinger4b779b32011-10-15 23:50:42 -0700276 result = user_function(*args, **kwds)
277 with lock:
Raymond Hettinger34d94a22012-04-30 14:14:28 -0700278 if key in cache:
Raymond Hettingerf96b2b02013-03-08 21:11:55 -0700279 # Getting here means that this same key was added to the
280 # cache while the lock was released. Since the link
Raymond Hettinger34d94a22012-04-30 14:14:28 -0700281 # update is already done, we need only return the
282 # computed result and update the count of misses.
283 pass
Raymond Hettinger018b4fb2012-04-30 20:48:55 -0700284 elif full:
Raymond Hettingerf96b2b02013-03-08 21:11:55 -0700285 # Use the old root to store the new key and result.
Raymond Hettingerf2c17a92013-03-04 03:34:09 -0500286 oldroot = root
287 oldroot[KEY] = key
288 oldroot[RESULT] = result
Raymond Hettingerf96b2b02013-03-08 21:11:55 -0700289 # Empty the oldest link and make it the new root.
290 # Keep a reference to the old key and old result to
291 # prevent their ref counts from going to zero during the
292 # update. That will prevent potentially arbitrary object
293 # clean-up code (i.e. __del__) from running while we're
294 # still adjusting the links.
Raymond Hettingerf2c17a92013-03-04 03:34:09 -0500295 root = oldroot[NEXT]
296 oldkey = root[KEY]
Raymond Hettingerf96b2b02013-03-08 21:11:55 -0700297 oldresult = root[RESULT]
Raymond Hettingerc6897852012-03-31 02:19:06 -0700298 root[KEY] = root[RESULT] = None
Raymond Hettingerf96b2b02013-03-08 21:11:55 -0700299 # Now update the cache dictionary.
Raymond Hettingerf2c17a92013-03-04 03:34:09 -0500300 del cache[oldkey]
Raymond Hettingerf96b2b02013-03-08 21:11:55 -0700301 # Save the potentially reentrant cache[key] assignment
302 # for last, after the root and links have been put in
303 # a consistent state.
Raymond Hettingerf2c17a92013-03-04 03:34:09 -0500304 cache[key] = oldroot
Raymond Hettinger018b4fb2012-04-30 20:48:55 -0700305 else:
Raymond Hettingerf96b2b02013-03-08 21:11:55 -0700306 # Put result in a new link at the front of the queue.
Raymond Hettinger018b4fb2012-04-30 20:48:55 -0700307 last = root[PREV]
308 link = [last, root, key, result]
Raymond Hettingerf2c17a92013-03-04 03:34:09 -0500309 last[NEXT] = root[PREV] = cache[key] = link
Raymond Hettinger018b4fb2012-04-30 20:48:55 -0700310 currsize += 1
Raymond Hettinger352cc8c2013-03-04 04:19:09 -0500311 full = (currsize >= maxsize)
Raymond Hettingerec0e9102012-03-16 01:16:31 -0700312 misses += 1
Raymond Hettingerc79fb0e2010-12-01 03:45:41 +0000313 return result
Georg Brandl2e7346a2010-07-31 18:09:23 +0000314
Nick Coghlan234515a2010-11-30 06:19:46 +0000315 def cache_info():
Raymond Hettinger5e20bab2010-11-30 07:13:04 +0000316 """Report cache statistics"""
Nick Coghlan234515a2010-11-30 06:19:46 +0000317 with lock:
Raymond Hettingerb6b98c02012-04-29 18:09:02 -0700318 return _CacheInfo(hits, misses, maxsize, currsize)
Nick Coghlan234515a2010-11-30 06:19:46 +0000319
Raymond Hettinger02566ec2010-09-04 22:46:06 +0000320 def cache_clear():
Benjamin Peterson1f594ad2010-08-08 13:17:07 +0000321 """Clear the cache and cache statistics"""
Raymond Hettinger018b4fb2012-04-30 20:48:55 -0700322 nonlocal hits, misses, currsize, full
Raymond Hettingercbe88132010-08-14 22:22:10 +0000323 with lock:
324 cache.clear()
Benjamin Peterson954cf572012-03-16 18:22:26 -0500325 root[:] = [root, root, None, None]
Raymond Hettingerb6b98c02012-04-29 18:09:02 -0700326 hits = misses = currsize = 0
Raymond Hettinger018b4fb2012-04-30 20:48:55 -0700327 full = False
Georg Brandl2e7346a2010-07-31 18:09:23 +0000328
Nick Coghlan234515a2010-11-30 06:19:46 +0000329 wrapper.cache_info = cache_info
Raymond Hettinger02566ec2010-09-04 22:46:06 +0000330 wrapper.cache_clear = cache_clear
Raymond Hettinger1ff50df2012-03-30 13:15:48 -0700331 return update_wrapper(wrapper, user_function)
Raymond Hettinger5fa40c02010-11-25 08:11:57 +0000332
Georg Brandl2e7346a2010-07-31 18:09:23 +0000333 return decorating_function