blob: 9d8cea9138cb628f4efca5e6eca6bdfbebe34dc1 [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',
12 'total_ordering', 'cmp_to_key', 'lfu_cache', 'lru_cache']
13
Guido van Rossum0919a1a2006-08-26 20:49:04 +000014from _functools import partial, reduce
Georg Brandl2e7346a2010-07-31 18:09:23 +000015from collections import OrderedDict, Counter
16from heapq import nsmallest
17from operator import itemgetter
Thomas Wouters4d70c3d2006-06-08 14:42:34 +000018
Thomas Wouters73e5a5b2006-06-08 15:35:45 +000019# update_wrapper() and wraps() are tools to help write
20# wrapper functions that can handle naive introspection
21
Antoine Pitrou560f7642010-08-04 18:28:02 +000022WRAPPER_ASSIGNMENTS = ('__module__', '__name__', '__doc__', '__annotations__')
Thomas Wouters73e5a5b2006-06-08 15:35:45 +000023WRAPPER_UPDATES = ('__dict__',)
24def update_wrapper(wrapper,
25 wrapped,
26 assigned = WRAPPER_ASSIGNMENTS,
27 updated = WRAPPER_UPDATES):
28 """Update a wrapper function to look like the wrapped function
29
30 wrapper is the function to be updated
31 wrapped is the original function
32 assigned is a tuple naming the attributes assigned directly
33 from the wrapped function to the wrapper function (defaults to
34 functools.WRAPPER_ASSIGNMENTS)
Thomas Wouters89f507f2006-12-13 04:49:30 +000035 updated is a tuple naming the attributes of the wrapper that
Thomas Wouters73e5a5b2006-06-08 15:35:45 +000036 are updated with the corresponding attribute from the wrapped
37 function (defaults to functools.WRAPPER_UPDATES)
38 """
39 for attr in assigned:
Antoine Pitrou560f7642010-08-04 18:28:02 +000040 if hasattr(wrapped, attr):
41 setattr(wrapper, attr, getattr(wrapped, attr))
Thomas Wouters73e5a5b2006-06-08 15:35:45 +000042 for attr in updated:
Thomas Wouters89f507f2006-12-13 04:49:30 +000043 getattr(wrapper, attr).update(getattr(wrapped, attr, {}))
Thomas Wouters73e5a5b2006-06-08 15:35:45 +000044 # Return the wrapper so this can be used as a decorator via partial()
45 return wrapper
46
47def wraps(wrapped,
48 assigned = WRAPPER_ASSIGNMENTS,
49 updated = WRAPPER_UPDATES):
50 """Decorator factory to apply update_wrapper() to a wrapper function
51
52 Returns a decorator that invokes update_wrapper() with the decorated
53 function as the wrapper argument and the arguments to wraps() as the
54 remaining arguments. Default arguments are as for update_wrapper().
55 This is a convenience function to simplify applying partial() to
56 update_wrapper().
57 """
58 return partial(update_wrapper, wrapped=wrapped,
59 assigned=assigned, updated=updated)
Raymond Hettingerc50846a2010-04-05 18:56:31 +000060
61def total_ordering(cls):
Georg Brandle5a26732010-05-19 21:06:36 +000062 """Class decorator that fills in missing ordering methods"""
Raymond Hettingerc50846a2010-04-05 18:56:31 +000063 convert = {
64 '__lt__': [('__gt__', lambda self, other: other < self),
65 ('__le__', lambda self, other: not other < self),
66 ('__ge__', lambda self, other: not self < other)],
67 '__le__': [('__ge__', lambda self, other: other <= self),
68 ('__lt__', lambda self, other: not other <= self),
69 ('__gt__', lambda self, other: not self <= other)],
70 '__gt__': [('__lt__', lambda self, other: other > self),
71 ('__ge__', lambda self, other: not other > self),
72 ('__le__', lambda self, other: not self > other)],
73 '__ge__': [('__le__', lambda self, other: other >= self),
74 ('__gt__', lambda self, other: not other >= self),
75 ('__lt__', lambda self, other: not self >= other)]
76 }
77 roots = set(dir(cls)) & set(convert)
Raymond Hettinger56de7e22010-04-10 16:59:03 +000078 if not roots:
79 raise ValueError('must define at least one ordering operation: < > <= >=')
80 root = max(roots) # prefer __lt__ to __le__ to __gt__ to __ge__
Raymond Hettingerc50846a2010-04-05 18:56:31 +000081 for opname, opfunc in convert[root]:
82 if opname not in roots:
83 opfunc.__name__ = opname
84 opfunc.__doc__ = getattr(int, opname).__doc__
85 setattr(cls, opname, opfunc)
86 return cls
87
88def cmp_to_key(mycmp):
Georg Brandle5a26732010-05-19 21:06:36 +000089 """Convert a cmp= function into a key= function"""
Raymond Hettingerc50846a2010-04-05 18:56:31 +000090 class K(object):
91 def __init__(self, obj, *args):
92 self.obj = obj
93 def __lt__(self, other):
94 return mycmp(self.obj, other.obj) < 0
95 def __gt__(self, other):
96 return mycmp(self.obj, other.obj) > 0
97 def __eq__(self, other):
98 return mycmp(self.obj, other.obj) == 0
99 def __le__(self, other):
100 return mycmp(self.obj, other.obj) <= 0
101 def __ge__(self, other):
102 return mycmp(self.obj, other.obj) >= 0
103 def __ne__(self, other):
104 return mycmp(self.obj, other.obj) != 0
105 def __hash__(self):
106 raise TypeError('hash not implemented')
107 return K
Georg Brandl2e7346a2010-07-31 18:09:23 +0000108
109def lfu_cache(maxsize=100):
Benjamin Peterson1f594ad2010-08-08 13:17:07 +0000110 """Least-frequently-used cache decorator.
Georg Brandl2e7346a2010-07-31 18:09:23 +0000111
112 Arguments to the cached function must be hashable.
113 Cache performance statistics stored in f.hits and f.misses.
114 Clear the cache using f.clear().
115 http://en.wikipedia.org/wiki/Cache_algorithms#Least-Frequently_Used
116
Benjamin Peterson1f594ad2010-08-08 13:17:07 +0000117 """
Georg Brandl2e7346a2010-07-31 18:09:23 +0000118 def decorating_function(user_function):
119 cache = {} # mapping of args to results
120 use_count = Counter() # times each key has been accessed
121 kwd_mark = object() # separate positional and keyword args
122
123 @wraps(user_function)
124 def wrapper(*args, **kwds):
125 key = args
126 if kwds:
127 key += (kwd_mark,) + tuple(sorted(kwds.items()))
128 use_count[key] += 1 # count a use of this key
129 try:
130 result = cache[key]
131 wrapper.hits += 1
132 except KeyError:
133 result = user_function(*args, **kwds)
134 cache[key] = result
135 wrapper.misses += 1
136 if len(cache) > maxsize:
137 # purge the 10% least frequently used entries
138 for key, _ in nsmallest(maxsize // 10,
139 use_count.items(),
140 key=itemgetter(1)):
141 del cache[key], use_count[key]
142 return result
143
144 def clear():
Benjamin Peterson1f594ad2010-08-08 13:17:07 +0000145 """Clear the cache and cache statistics"""
Georg Brandl2e7346a2010-07-31 18:09:23 +0000146 cache.clear()
147 use_count.clear()
148 wrapper.hits = wrapper.misses = 0
149
150 wrapper.hits = wrapper.misses = 0
151 wrapper.clear = clear
152 return wrapper
153 return decorating_function
154
155def lru_cache(maxsize=100):
Benjamin Peterson1f594ad2010-08-08 13:17:07 +0000156 """Least-recently-used cache decorator.
Georg Brandl2e7346a2010-07-31 18:09:23 +0000157
158 Arguments to the cached function must be hashable.
159 Cache performance statistics stored in f.hits and f.misses.
160 Clear the cache using f.clear().
161 http://en.wikipedia.org/wiki/Cache_algorithms#Least_Recently_Used
162
Benjamin Peterson1f594ad2010-08-08 13:17:07 +0000163 """
Georg Brandl2e7346a2010-07-31 18:09:23 +0000164 def decorating_function(user_function):
165 cache = OrderedDict() # ordered least recent to most recent
166 kwd_mark = object() # separate positional and keyword args
167
168 @wraps(user_function)
169 def wrapper(*args, **kwds):
170 key = args
171 if kwds:
172 key += (kwd_mark,) + tuple(sorted(kwds.items()))
173 try:
174 result = cache.pop(key)
175 wrapper.hits += 1
176 except KeyError:
177 result = user_function(*args, **kwds)
178 wrapper.misses += 1
179 if len(cache) >= maxsize:
180 cache.popitem(0) # purge least recently used cache entry
181 cache[key] = result # record recent use of this key
182 return result
183
184 def clear():
Benjamin Peterson1f594ad2010-08-08 13:17:07 +0000185 """Clear the cache and cache statistics"""
Georg Brandl2e7346a2010-07-31 18:09:23 +0000186 cache.clear()
187 wrapper.hits = wrapper.misses = 0
188
189 wrapper.hits = wrapper.misses = 0
190 wrapper.clear = clear
191 return wrapper
192 return decorating_function