blob: fa11d2f27c05f4604efe724944decde992809aac [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
Antoine Pitroub5b37142012-11-13 21:35:40 +010014try:
15 from _functools import reduce
16except ImportError:
17 pass
Raymond Hettingerec0e9102012-03-16 01:16:31 -070018from collections import namedtuple
Raymond Hettingercbe88132010-08-14 22:22:10 +000019try:
Raymond Hettingerfd541172013-03-01 03:47:57 -080020 from _thread import RLock
Raymond Hettingercbe88132010-08-14 22:22:10 +000021except:
Raymond Hettinger409f6632013-03-01 23:20:13 -080022 class RLock:
Raymond Hettingerf96b2b02013-03-08 21:11:55 -070023 'Dummy reentrant lock for builds without threads'
Raymond Hettinger409f6632013-03-01 23:20:13 -080024 def __enter__(self): pass
25 def __exit__(self, exctype, excinst, exctb): pass
Thomas Wouters4d70c3d2006-06-08 14:42:34 +000026
Raymond Hettingerbc8e81d2012-03-17 00:24:09 -070027
28################################################################################
29### update_wrapper() and wraps() decorator
30################################################################################
31
Thomas Wouters73e5a5b2006-06-08 15:35:45 +000032# update_wrapper() and wraps() are tools to help write
33# wrapper functions that can handle naive introspection
34
Meador Ingeff7f64c2011-12-11 22:37:31 -060035WRAPPER_ASSIGNMENTS = ('__module__', '__name__', '__qualname__', '__doc__',
36 '__annotations__')
Thomas Wouters73e5a5b2006-06-08 15:35:45 +000037WRAPPER_UPDATES = ('__dict__',)
38def update_wrapper(wrapper,
39 wrapped,
40 assigned = WRAPPER_ASSIGNMENTS,
41 updated = WRAPPER_UPDATES):
42 """Update a wrapper function to look like the wrapped function
43
44 wrapper is the function to be updated
45 wrapped is the original function
46 assigned is a tuple naming the attributes assigned directly
47 from the wrapped function to the wrapper function (defaults to
48 functools.WRAPPER_ASSIGNMENTS)
Thomas Wouters89f507f2006-12-13 04:49:30 +000049 updated is a tuple naming the attributes of the wrapper that
Thomas Wouters73e5a5b2006-06-08 15:35:45 +000050 are updated with the corresponding attribute from the wrapped
51 function (defaults to functools.WRAPPER_UPDATES)
52 """
Nick Coghlan98876832010-08-17 06:17:18 +000053 wrapper.__wrapped__ = wrapped
Thomas Wouters73e5a5b2006-06-08 15:35:45 +000054 for attr in assigned:
Nick Coghlan98876832010-08-17 06:17:18 +000055 try:
56 value = getattr(wrapped, attr)
57 except AttributeError:
58 pass
59 else:
60 setattr(wrapper, attr, value)
Thomas Wouters73e5a5b2006-06-08 15:35:45 +000061 for attr in updated:
Thomas Wouters89f507f2006-12-13 04:49:30 +000062 getattr(wrapper, attr).update(getattr(wrapped, attr, {}))
Thomas Wouters73e5a5b2006-06-08 15:35:45 +000063 # Return the wrapper so this can be used as a decorator via partial()
64 return wrapper
65
66def wraps(wrapped,
67 assigned = WRAPPER_ASSIGNMENTS,
68 updated = WRAPPER_UPDATES):
69 """Decorator factory to apply update_wrapper() to a wrapper function
70
71 Returns a decorator that invokes update_wrapper() with the decorated
72 function as the wrapper argument and the arguments to wraps() as the
73 remaining arguments. Default arguments are as for update_wrapper().
74 This is a convenience function to simplify applying partial() to
75 update_wrapper().
76 """
77 return partial(update_wrapper, wrapped=wrapped,
78 assigned=assigned, updated=updated)
Raymond Hettingerc50846a2010-04-05 18:56:31 +000079
Raymond Hettingerbc8e81d2012-03-17 00:24:09 -070080
81################################################################################
82### total_ordering class decorator
83################################################################################
84
Raymond Hettingerc50846a2010-04-05 18:56:31 +000085def total_ordering(cls):
Georg Brandle5a26732010-05-19 21:06:36 +000086 """Class decorator that fills in missing ordering methods"""
Raymond Hettingerc50846a2010-04-05 18:56:31 +000087 convert = {
Raymond Hettinger23f9fc32011-01-08 07:01:56 +000088 '__lt__': [('__gt__', lambda self, other: not (self < other or self == other)),
89 ('__le__', lambda self, other: self < other or self == other),
Raymond Hettingerc50846a2010-04-05 18:56:31 +000090 ('__ge__', lambda self, other: not self < other)],
Raymond Hettinger23f9fc32011-01-08 07:01:56 +000091 '__le__': [('__ge__', lambda self, other: not self <= other or self == other),
92 ('__lt__', lambda self, other: self <= other and not self == other),
Raymond Hettingerc50846a2010-04-05 18:56:31 +000093 ('__gt__', lambda self, other: not self <= other)],
Raymond Hettinger23f9fc32011-01-08 07:01:56 +000094 '__gt__': [('__lt__', lambda self, other: not (self > other or self == other)),
95 ('__ge__', lambda self, other: self > other or self == other),
Raymond Hettingerc50846a2010-04-05 18:56:31 +000096 ('__le__', lambda self, other: not self > other)],
Raymond Hettinger23f9fc32011-01-08 07:01:56 +000097 '__ge__': [('__le__', lambda self, other: (not self >= other) or self == other),
98 ('__gt__', lambda self, other: self >= other and not self == other),
Raymond Hettingerc50846a2010-04-05 18:56:31 +000099 ('__lt__', lambda self, other: not self >= other)]
100 }
Raymond Hettinger3255c632010-09-16 00:31:21 +0000101 # Find user-defined comparisons (not those inherited from object).
Raymond Hettinger1006bd42010-09-14 22:55:13 +0000102 roots = [op for op in convert if getattr(cls, op, None) is not getattr(object, op, None)]
Raymond Hettinger56de7e22010-04-10 16:59:03 +0000103 if not roots:
104 raise ValueError('must define at least one ordering operation: < > <= >=')
105 root = max(roots) # prefer __lt__ to __le__ to __gt__ to __ge__
Raymond Hettingerc50846a2010-04-05 18:56:31 +0000106 for opname, opfunc in convert[root]:
107 if opname not in roots:
108 opfunc.__name__ = opname
109 opfunc.__doc__ = getattr(int, opname).__doc__
110 setattr(cls, opname, opfunc)
111 return cls
112
Raymond Hettingerbc8e81d2012-03-17 00:24:09 -0700113
114################################################################################
115### cmp_to_key() function converter
116################################################################################
117
Raymond Hettingerc50846a2010-04-05 18:56:31 +0000118def cmp_to_key(mycmp):
Georg Brandle5a26732010-05-19 21:06:36 +0000119 """Convert a cmp= function into a key= function"""
Raymond Hettingerc50846a2010-04-05 18:56:31 +0000120 class K(object):
Raymond Hettingera0d1d962011-03-21 17:50:28 -0700121 __slots__ = ['obj']
Raymond Hettinger7ab9e222011-04-05 02:33:54 -0700122 def __init__(self, obj):
Raymond Hettingerc50846a2010-04-05 18:56:31 +0000123 self.obj = obj
124 def __lt__(self, other):
125 return mycmp(self.obj, other.obj) < 0
126 def __gt__(self, other):
127 return mycmp(self.obj, other.obj) > 0
128 def __eq__(self, other):
129 return mycmp(self.obj, other.obj) == 0
130 def __le__(self, other):
131 return mycmp(self.obj, other.obj) <= 0
132 def __ge__(self, other):
133 return mycmp(self.obj, other.obj) >= 0
134 def __ne__(self, other):
135 return mycmp(self.obj, other.obj) != 0
Raymond Hettinger003be522011-05-03 11:01:32 -0700136 __hash__ = None
Raymond Hettingerc50846a2010-04-05 18:56:31 +0000137 return K
Georg Brandl2e7346a2010-07-31 18:09:23 +0000138
Raymond Hettinger7ab9e222011-04-05 02:33:54 -0700139try:
140 from _functools import cmp_to_key
141except ImportError:
142 pass
143
Raymond Hettingerbc8e81d2012-03-17 00:24:09 -0700144
145################################################################################
Antoine Pitroub5b37142012-11-13 21:35:40 +0100146### partial() argument application
147################################################################################
148
149def partial(func, *args, **keywords):
150 """new function with partial application of the given arguments
151 and keywords.
152 """
153 def newfunc(*fargs, **fkeywords):
154 newkeywords = keywords.copy()
155 newkeywords.update(fkeywords)
156 return func(*(args + fargs), **newkeywords)
157 newfunc.func = func
158 newfunc.args = args
159 newfunc.keywords = keywords
160 return newfunc
161
162try:
163 from _functools import partial
164except ImportError:
165 pass
166
167
168################################################################################
Raymond Hettingerbc8e81d2012-03-17 00:24:09 -0700169### LRU Cache function decorator
170################################################################################
171
Raymond Hettingerdce583e2012-03-16 22:12:20 -0700172_CacheInfo = namedtuple("CacheInfo", ["hits", "misses", "maxsize", "currsize"])
Nick Coghlan234515a2010-11-30 06:19:46 +0000173
Raymond Hettinger0c9050c2012-06-04 00:21:14 -0700174class _HashedSeq(list):
Raymond Hettingerf96b2b02013-03-08 21:11:55 -0700175 """ This class guarantees that hash() will be called no more than once
176 per element. This is important because the lru_cache() will hash
177 the key multiple times on a cache miss.
178
179 """
180
Raymond Hettinger9acbb602012-04-30 22:32:16 -0700181 __slots__ = 'hashvalue'
182
Raymond Hettinger0c9050c2012-06-04 00:21:14 -0700183 def __init__(self, tup, hash=hash):
184 self[:] = tup
185 self.hashvalue = hash(tup)
Raymond Hettinger9acbb602012-04-30 22:32:16 -0700186
187 def __hash__(self):
188 return self.hashvalue
189
Raymond Hettinger0c9050c2012-06-04 00:21:14 -0700190def _make_key(args, kwds, typed,
191 kwd_mark = (object(),),
192 fasttypes = {int, str, frozenset, type(None)},
193 sorted=sorted, tuple=tuple, type=type, len=len):
Raymond Hettingerf96b2b02013-03-08 21:11:55 -0700194 """Make a cache key from optionally typed positional and keyword arguments
195
196 The key is constructed in a way that is flat as possible rather than
197 as a nested structure that would take more memory.
198
199 If there is only a single argument and its data type is known to cache
200 its hash value, then that argument is returned without a wrapper. This
201 saves space and improves lookup speed.
202
203 """
Raymond Hettinger0c9050c2012-06-04 00:21:14 -0700204 key = args
205 if kwds:
206 sorted_items = sorted(kwds.items())
207 key += kwd_mark
208 for item in sorted_items:
209 key += item
210 if typed:
211 key += tuple(type(v) for v in args)
212 if kwds:
213 key += tuple(type(v) for k, v in sorted_items)
214 elif len(key) == 1 and type(key[0]) in fasttypes:
215 return key[0]
216 return _HashedSeq(key)
217
Raymond Hettinger010ce322012-05-19 21:20:48 -0700218def lru_cache(maxsize=128, typed=False):
Benjamin Peterson1f594ad2010-08-08 13:17:07 +0000219 """Least-recently-used cache decorator.
Georg Brandl2e7346a2010-07-31 18:09:23 +0000220
Raymond Hettingerc79fb0e2010-12-01 03:45:41 +0000221 If *maxsize* is set to None, the LRU features are disabled and the cache
222 can grow without bound.
223
Raymond Hettingercd9fdfd2011-10-20 08:57:45 -0700224 If *typed* is True, arguments of different types will be cached separately.
225 For example, f(3.0) and f(3) will be treated as distinct calls with
226 distinct results.
227
Georg Brandl2e7346a2010-07-31 18:09:23 +0000228 Arguments to the cached function must be hashable.
Raymond Hettinger5fa40c02010-11-25 08:11:57 +0000229
Raymond Hettinger7f7a5a72012-03-30 21:50:40 -0700230 View the cache statistics named tuple (hits, misses, maxsize, currsize)
231 with f.cache_info(). Clear the cache and statistics with f.cache_clear().
Raymond Hettinger00f2f972010-12-01 00:47:56 +0000232 Access the underlying function with f.__wrapped__.
Raymond Hettinger5fa40c02010-11-25 08:11:57 +0000233
234 See: http://en.wikipedia.org/wiki/Cache_algorithms#Least_Recently_Used
Georg Brandl2e7346a2010-07-31 18:09:23 +0000235
Benjamin Peterson1f594ad2010-08-08 13:17:07 +0000236 """
Raymond Hettinger1ff50df2012-03-30 13:15:48 -0700237
Raymond Hettinger5fa40c02010-11-25 08:11:57 +0000238 # Users should only access the lru_cache through its public API:
Raymond Hettinger5e20bab2010-11-30 07:13:04 +0000239 # cache_info, cache_clear, and f.__wrapped__
Raymond Hettinger5fa40c02010-11-25 08:11:57 +0000240 # The internals of the lru_cache are encapsulated for thread safety and
241 # to allow the implementation to change (including a possible C version).
242
Raymond Hettinger9f0ab9f2012-04-29 14:55:27 -0700243 # Constants shared by all lru cache instances:
Raymond Hettingerb6b98c02012-04-29 18:09:02 -0700244 sentinel = object() # unique object used to signal cache misses
Raymond Hettinger0c9050c2012-06-04 00:21:14 -0700245 make_key = _make_key # build a key from the function arguments
Raymond Hettinger9f0ab9f2012-04-29 14:55:27 -0700246 PREV, NEXT, KEY, RESULT = 0, 1, 2, 3 # names for the link fields
247
Raymond Hettinger6e8c8172012-03-16 16:53:05 -0700248 def decorating_function(user_function):
Raymond Hettinger7f7a5a72012-03-30 21:50:40 -0700249 cache = {}
Raymond Hettinger832edde2013-02-17 00:08:45 -0800250 hits = misses = 0
Raymond Hettinger018b4fb2012-04-30 20:48:55 -0700251 full = False
Raymond Hettingerc6897852012-03-31 02:19:06 -0700252 cache_get = cache.get # bound method to lookup a key or return None
Raymond Hettingerfd541172013-03-01 03:47:57 -0800253 lock = RLock() # because linkedlist updates aren't threadsafe
Raymond Hettinger7f7a5a72012-03-30 21:50:40 -0700254 root = [] # root of the circular doubly linked list
255 root[:] = [root, root, None, None] # initialize by pointing to self
Raymond Hettinger6e8c8172012-03-16 16:53:05 -0700256
Raymond Hettinger7e0c5812012-03-17 15:10:24 -0700257 if maxsize == 0:
258
Raymond Hettinger7e0c5812012-03-17 15:10:24 -0700259 def wrapper(*args, **kwds):
Raymond Hettingerf96b2b02013-03-08 21:11:55 -0700260 # No caching -- just a statistics update after a successful call
Raymond Hettinger7e0c5812012-03-17 15:10:24 -0700261 nonlocal misses
Raymond Hettinger7dabfed2012-03-17 15:11:09 -0700262 result = user_function(*args, **kwds)
Raymond Hettinger7e0c5812012-03-17 15:10:24 -0700263 misses += 1
264 return result
265
266 elif maxsize is None:
Raymond Hettingerbc8e81d2012-03-17 00:24:09 -0700267
Raymond Hettingerc79fb0e2010-12-01 03:45:41 +0000268 def wrapper(*args, **kwds):
Raymond Hettingerf96b2b02013-03-08 21:11:55 -0700269 # Simple caching without ordering or size limit
Raymond Hettinger832edde2013-02-17 00:08:45 -0800270 nonlocal hits, misses
Raymond Hettinger9acbb602012-04-30 22:32:16 -0700271 key = make_key(args, kwds, typed)
Raymond Hettinger7f7a5a72012-03-30 21:50:40 -0700272 result = cache_get(key, sentinel)
273 if result is not sentinel:
Nick Coghlan234515a2010-11-30 06:19:46 +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 cache[key] = result
278 misses += 1
Raymond Hettingerc79fb0e2010-12-01 03:45:41 +0000279 return result
Raymond Hettingerbc8e81d2012-03-17 00:24:09 -0700280
Raymond Hettingerc79fb0e2010-12-01 03:45:41 +0000281 else:
Raymond Hettingerbc8e81d2012-03-17 00:24:09 -0700282
Raymond Hettingerc79fb0e2010-12-01 03:45:41 +0000283 def wrapper(*args, **kwds):
Raymond Hettingerf96b2b02013-03-08 21:11:55 -0700284 # Size limited caching that tracks accesses by recency
Raymond Hettinger832edde2013-02-17 00:08:45 -0800285 nonlocal root, hits, misses, full
Raymond Hettinger9acbb602012-04-30 22:32:16 -0700286 key = make_key(args, kwds, typed)
Raymond Hettinger4b779b32011-10-15 23:50:42 -0700287 with lock:
Raymond Hettingerec0e9102012-03-16 01:16:31 -0700288 link = cache_get(key)
289 if link is not None:
Raymond Hettingerf96b2b02013-03-08 21:11:55 -0700290 # Move the link to the front of the circular queue
291 link_prev, link_next, _key, result = link
Raymond Hettingerec0e9102012-03-16 01:16:31 -0700292 link_prev[NEXT] = link_next
293 link_next[PREV] = link_prev
294 last = root[PREV]
295 last[NEXT] = root[PREV] = link
296 link[PREV] = last
297 link[NEXT] = root
Raymond Hettingerc79fb0e2010-12-01 03:45:41 +0000298 hits += 1
Raymond Hettinger4b779b32011-10-15 23:50:42 -0700299 return result
Raymond Hettinger4b779b32011-10-15 23:50:42 -0700300 result = user_function(*args, **kwds)
301 with lock:
Raymond Hettinger34d94a22012-04-30 14:14:28 -0700302 if key in cache:
Raymond Hettingerf96b2b02013-03-08 21:11:55 -0700303 # Getting here means that this same key was added to the
304 # cache while the lock was released. Since the link
Raymond Hettinger34d94a22012-04-30 14:14:28 -0700305 # update is already done, we need only return the
306 # computed result and update the count of misses.
307 pass
Raymond Hettinger018b4fb2012-04-30 20:48:55 -0700308 elif full:
Raymond Hettingerf96b2b02013-03-08 21:11:55 -0700309 # Use the old root to store the new key and result.
Raymond Hettingerf2c17a92013-03-04 03:34:09 -0500310 oldroot = root
311 oldroot[KEY] = key
312 oldroot[RESULT] = result
Raymond Hettingerf96b2b02013-03-08 21:11:55 -0700313 # Empty the oldest link and make it the new root.
314 # Keep a reference to the old key and old result to
315 # prevent their ref counts from going to zero during the
316 # update. That will prevent potentially arbitrary object
317 # clean-up code (i.e. __del__) from running while we're
318 # still adjusting the links.
Raymond Hettingerf2c17a92013-03-04 03:34:09 -0500319 root = oldroot[NEXT]
320 oldkey = root[KEY]
Raymond Hettingerf96b2b02013-03-08 21:11:55 -0700321 oldresult = root[RESULT]
Raymond Hettingerc6897852012-03-31 02:19:06 -0700322 root[KEY] = root[RESULT] = None
Raymond Hettingerf96b2b02013-03-08 21:11:55 -0700323 # Now update the cache dictionary.
Raymond Hettingerf2c17a92013-03-04 03:34:09 -0500324 del cache[oldkey]
Raymond Hettingerf96b2b02013-03-08 21:11:55 -0700325 # Save the potentially reentrant cache[key] assignment
326 # for last, after the root and links have been put in
327 # a consistent state.
Raymond Hettingerf2c17a92013-03-04 03:34:09 -0500328 cache[key] = oldroot
Raymond Hettinger018b4fb2012-04-30 20:48:55 -0700329 else:
Raymond Hettingerf96b2b02013-03-08 21:11:55 -0700330 # Put result in a new link at the front of the queue.
Raymond Hettinger018b4fb2012-04-30 20:48:55 -0700331 last = root[PREV]
332 link = [last, root, key, result]
Raymond Hettingerf2c17a92013-03-04 03:34:09 -0500333 last[NEXT] = root[PREV] = cache[key] = link
Raymond Hettingerbb5f4802013-03-04 04:20:46 -0500334 full = (len(cache) >= maxsize)
Raymond Hettingerec0e9102012-03-16 01:16:31 -0700335 misses += 1
Raymond Hettingerc79fb0e2010-12-01 03:45:41 +0000336 return result
Georg Brandl2e7346a2010-07-31 18:09:23 +0000337
Nick Coghlan234515a2010-11-30 06:19:46 +0000338 def cache_info():
Raymond Hettinger5e20bab2010-11-30 07:13:04 +0000339 """Report cache statistics"""
Nick Coghlan234515a2010-11-30 06:19:46 +0000340 with lock:
Raymond Hettinger832edde2013-02-17 00:08:45 -0800341 return _CacheInfo(hits, misses, maxsize, len(cache))
Nick Coghlan234515a2010-11-30 06:19:46 +0000342
Raymond Hettinger02566ec2010-09-04 22:46:06 +0000343 def cache_clear():
Benjamin Peterson1f594ad2010-08-08 13:17:07 +0000344 """Clear the cache and cache statistics"""
Raymond Hettinger832edde2013-02-17 00:08:45 -0800345 nonlocal hits, misses, full
Raymond Hettingercbe88132010-08-14 22:22:10 +0000346 with lock:
347 cache.clear()
Benjamin Peterson954cf572012-03-16 18:22:26 -0500348 root[:] = [root, root, None, None]
Raymond Hettinger832edde2013-02-17 00:08:45 -0800349 hits = misses = 0
Raymond Hettinger018b4fb2012-04-30 20:48:55 -0700350 full = False
Georg Brandl2e7346a2010-07-31 18:09:23 +0000351
Nick Coghlan234515a2010-11-30 06:19:46 +0000352 wrapper.cache_info = cache_info
Raymond Hettinger02566ec2010-09-04 22:46:06 +0000353 wrapper.cache_clear = cache_clear
Raymond Hettinger1ff50df2012-03-30 13:15:48 -0700354 return update_wrapper(wrapper, user_function)
Raymond Hettinger5fa40c02010-11-25 08:11:57 +0000355
Georg Brandl2e7346a2010-07-31 18:09:23 +0000356 return decorating_function