blob: d4506344859e311e7eddceed17111fb6965fd051 [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
Nick Coghlan234515a2010-11-30 06:19:46 +000015from collections import OrderedDict, namedtuple
Raymond Hettingercbe88132010-08-14 22:22:10 +000016try:
17 from _thread import allocate_lock as Lock
18except:
19 from _dummy_thread import allocate_lock as Lock
Thomas Wouters4d70c3d2006-06-08 14:42:34 +000020
Thomas Wouters73e5a5b2006-06-08 15:35:45 +000021# update_wrapper() and wraps() are tools to help write
22# wrapper functions that can handle naive introspection
23
Antoine Pitrou560f7642010-08-04 18:28:02 +000024WRAPPER_ASSIGNMENTS = ('__module__', '__name__', '__doc__', '__annotations__')
Thomas Wouters73e5a5b2006-06-08 15:35:45 +000025WRAPPER_UPDATES = ('__dict__',)
26def update_wrapper(wrapper,
27 wrapped,
28 assigned = WRAPPER_ASSIGNMENTS,
29 updated = WRAPPER_UPDATES):
30 """Update a wrapper function to look like the wrapped function
31
32 wrapper is the function to be updated
33 wrapped is the original function
34 assigned is a tuple naming the attributes assigned directly
35 from the wrapped function to the wrapper function (defaults to
36 functools.WRAPPER_ASSIGNMENTS)
Thomas Wouters89f507f2006-12-13 04:49:30 +000037 updated is a tuple naming the attributes of the wrapper that
Thomas Wouters73e5a5b2006-06-08 15:35:45 +000038 are updated with the corresponding attribute from the wrapped
39 function (defaults to functools.WRAPPER_UPDATES)
40 """
Nick Coghlan98876832010-08-17 06:17:18 +000041 wrapper.__wrapped__ = wrapped
Thomas Wouters73e5a5b2006-06-08 15:35:45 +000042 for attr in assigned:
Nick Coghlan98876832010-08-17 06:17:18 +000043 try:
44 value = getattr(wrapped, attr)
45 except AttributeError:
46 pass
47 else:
48 setattr(wrapper, attr, value)
Thomas Wouters73e5a5b2006-06-08 15:35:45 +000049 for attr in updated:
Thomas Wouters89f507f2006-12-13 04:49:30 +000050 getattr(wrapper, attr).update(getattr(wrapped, attr, {}))
Thomas Wouters73e5a5b2006-06-08 15:35:45 +000051 # Return the wrapper so this can be used as a decorator via partial()
52 return wrapper
53
54def wraps(wrapped,
55 assigned = WRAPPER_ASSIGNMENTS,
56 updated = WRAPPER_UPDATES):
57 """Decorator factory to apply update_wrapper() to a wrapper function
58
59 Returns a decorator that invokes update_wrapper() with the decorated
60 function as the wrapper argument and the arguments to wraps() as the
61 remaining arguments. Default arguments are as for update_wrapper().
62 This is a convenience function to simplify applying partial() to
63 update_wrapper().
64 """
65 return partial(update_wrapper, wrapped=wrapped,
66 assigned=assigned, updated=updated)
Raymond Hettingerc50846a2010-04-05 18:56:31 +000067
68def total_ordering(cls):
Georg Brandle5a26732010-05-19 21:06:36 +000069 """Class decorator that fills in missing ordering methods"""
Raymond Hettingerc50846a2010-04-05 18:56:31 +000070 convert = {
71 '__lt__': [('__gt__', lambda self, other: other < self),
72 ('__le__', lambda self, other: not other < self),
73 ('__ge__', lambda self, other: not self < other)],
74 '__le__': [('__ge__', lambda self, other: other <= self),
75 ('__lt__', lambda self, other: not other <= self),
76 ('__gt__', lambda self, other: not self <= other)],
77 '__gt__': [('__lt__', lambda self, other: other > self),
78 ('__ge__', lambda self, other: not other > self),
79 ('__le__', lambda self, other: not self > other)],
80 '__ge__': [('__le__', lambda self, other: other >= self),
81 ('__gt__', lambda self, other: not other >= self),
82 ('__lt__', lambda self, other: not self >= other)]
83 }
Raymond Hettinger3255c632010-09-16 00:31:21 +000084 # Find user-defined comparisons (not those inherited from object).
Raymond Hettinger1006bd42010-09-14 22:55:13 +000085 roots = [op for op in convert if getattr(cls, op, None) is not getattr(object, op, None)]
Raymond Hettinger56de7e22010-04-10 16:59:03 +000086 if not roots:
87 raise ValueError('must define at least one ordering operation: < > <= >=')
88 root = max(roots) # prefer __lt__ to __le__ to __gt__ to __ge__
Raymond Hettingerc50846a2010-04-05 18:56:31 +000089 for opname, opfunc in convert[root]:
90 if opname not in roots:
91 opfunc.__name__ = opname
92 opfunc.__doc__ = getattr(int, opname).__doc__
93 setattr(cls, opname, opfunc)
94 return cls
95
96def cmp_to_key(mycmp):
Georg Brandle5a26732010-05-19 21:06:36 +000097 """Convert a cmp= function into a key= function"""
Raymond Hettingerc50846a2010-04-05 18:56:31 +000098 class K(object):
99 def __init__(self, obj, *args):
100 self.obj = obj
101 def __lt__(self, other):
102 return mycmp(self.obj, other.obj) < 0
103 def __gt__(self, other):
104 return mycmp(self.obj, other.obj) > 0
105 def __eq__(self, other):
106 return mycmp(self.obj, other.obj) == 0
107 def __le__(self, other):
108 return mycmp(self.obj, other.obj) <= 0
109 def __ge__(self, other):
110 return mycmp(self.obj, other.obj) >= 0
111 def __ne__(self, other):
112 return mycmp(self.obj, other.obj) != 0
113 def __hash__(self):
114 raise TypeError('hash not implemented')
115 return K
Georg Brandl2e7346a2010-07-31 18:09:23 +0000116
Raymond Hettinger7496b412010-11-30 19:15:45 +0000117_CacheInfo = namedtuple("CacheInfo", "hits misses maxsize currsize")
Nick Coghlan234515a2010-11-30 06:19:46 +0000118
Georg Brandl2e7346a2010-07-31 18:09:23 +0000119def lru_cache(maxsize=100):
Benjamin Peterson1f594ad2010-08-08 13:17:07 +0000120 """Least-recently-used cache decorator.
Georg Brandl2e7346a2010-07-31 18:09:23 +0000121
Raymond Hettingerc79fb0e2010-12-01 03:45:41 +0000122 If *maxsize* is set to None, the LRU features are disabled and the cache
123 can grow without bound.
124
Georg Brandl2e7346a2010-07-31 18:09:23 +0000125 Arguments to the cached function must be hashable.
Raymond Hettinger5fa40c02010-11-25 08:11:57 +0000126
Raymond Hettinger00f2f972010-12-01 00:47:56 +0000127 View the cache statistics named tuple (hits, misses, maxsize, currsize) with
Raymond Hettinger5e20bab2010-11-30 07:13:04 +0000128 f.cache_info(). Clear the cache and statistics with f.cache_clear().
Raymond Hettinger00f2f972010-12-01 00:47:56 +0000129 Access the underlying function with f.__wrapped__.
Raymond Hettinger5fa40c02010-11-25 08:11:57 +0000130
131 See: http://en.wikipedia.org/wiki/Cache_algorithms#Least_Recently_Used
Georg Brandl2e7346a2010-07-31 18:09:23 +0000132
Benjamin Peterson1f594ad2010-08-08 13:17:07 +0000133 """
Raymond Hettinger5fa40c02010-11-25 08:11:57 +0000134 # Users should only access the lru_cache through its public API:
Raymond Hettinger5e20bab2010-11-30 07:13:04 +0000135 # cache_info, cache_clear, and f.__wrapped__
Raymond Hettinger5fa40c02010-11-25 08:11:57 +0000136 # The internals of the lru_cache are encapsulated for thread safety and
137 # to allow the implementation to change (including a possible C version).
138
139 def decorating_function(user_function,
140 tuple=tuple, sorted=sorted, len=len, KeyError=KeyError):
141
Nick Coghlan234515a2010-11-30 06:19:46 +0000142 hits = misses = 0
Raymond Hettinger5e20bab2010-11-30 07:13:04 +0000143 kwd_mark = object() # separates positional and keyword args
Raymond Hettingercbe88132010-08-14 22:22:10 +0000144 lock = Lock()
Georg Brandl2e7346a2010-07-31 18:09:23 +0000145
Raymond Hettingerc79fb0e2010-12-01 03:45:41 +0000146 if maxsize is None:
147 cache = dict() # simple cache without ordering or size limit
148
149 @wraps(user_function)
150 def wrapper(*args, **kwds):
151 nonlocal hits, misses
152 key = args
153 if kwds:
154 key += (kwd_mark,) + tuple(sorted(kwds.items()))
155 try:
Raymond Hettingercbe88132010-08-14 22:22:10 +0000156 result = cache[key]
Nick Coghlan234515a2010-11-30 06:19:46 +0000157 hits += 1
Raymond Hettingerc79fb0e2010-12-01 03:45:41 +0000158 except KeyError:
159 result = user_function(*args, **kwds)
160 cache[key] = result
Nick Coghlan234515a2010-11-30 06:19:46 +0000161 misses += 1
Raymond Hettingerc79fb0e2010-12-01 03:45:41 +0000162 return result
163 else:
164 cache = OrderedDict() # ordered least recent to most recent
165 cache_popitem = cache.popitem
166 cache_renew = cache.move_to_end
167
168 @wraps(user_function)
169 def wrapper(*args, **kwds):
170 nonlocal hits, misses
171 key = args
172 if kwds:
173 key += (kwd_mark,) + tuple(sorted(kwds.items()))
174 try:
175 with lock:
176 result = cache[key]
177 cache_renew(key) # record recent use of this key
178 hits += 1
179 except KeyError:
180 result = user_function(*args, **kwds)
181 with lock:
182 cache[key] = result # record recent use of this key
183 misses += 1
184 if len(cache) > maxsize:
185 cache_popitem(0) # purge least recently used cache entry
186 return result
Georg Brandl2e7346a2010-07-31 18:09:23 +0000187
Nick Coghlan234515a2010-11-30 06:19:46 +0000188 def cache_info():
Raymond Hettinger5e20bab2010-11-30 07:13:04 +0000189 """Report cache statistics"""
Nick Coghlan234515a2010-11-30 06:19:46 +0000190 with lock:
Raymond Hettinger7496b412010-11-30 19:15:45 +0000191 return _CacheInfo(hits, misses, maxsize, len(cache))
Nick Coghlan234515a2010-11-30 06:19:46 +0000192
Raymond Hettinger02566ec2010-09-04 22:46:06 +0000193 def cache_clear():
Benjamin Peterson1f594ad2010-08-08 13:17:07 +0000194 """Clear the cache and cache statistics"""
Nick Coghlan234515a2010-11-30 06:19:46 +0000195 nonlocal hits, misses
Raymond Hettingercbe88132010-08-14 22:22:10 +0000196 with lock:
197 cache.clear()
Nick Coghlan234515a2010-11-30 06:19:46 +0000198 hits = misses = 0
Georg Brandl2e7346a2010-07-31 18:09:23 +0000199
Nick Coghlan234515a2010-11-30 06:19:46 +0000200 wrapper.cache_info = cache_info
Raymond Hettinger02566ec2010-09-04 22:46:06 +0000201 wrapper.cache_clear = cache_clear
Georg Brandl2e7346a2010-07-31 18:09:23 +0000202 return wrapper
Raymond Hettinger5fa40c02010-11-25 08:11:57 +0000203
Georg Brandl2e7346a2010-07-31 18:09:23 +0000204 return decorating_function