blob: efc533cf7d1a47f0ec0bdbb834b1b32bec8961b2 [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
Raymond Hettingercbe88132010-08-14 22:22:10 +000018try:
19 from _thread import allocate_lock as Lock
20except:
21 from _dummy_thread import allocate_lock as Lock
Thomas Wouters4d70c3d2006-06-08 14:42:34 +000022
Thomas Wouters73e5a5b2006-06-08 15:35:45 +000023# update_wrapper() and wraps() are tools to help write
24# wrapper functions that can handle naive introspection
25
Antoine Pitrou560f7642010-08-04 18:28:02 +000026WRAPPER_ASSIGNMENTS = ('__module__', '__name__', '__doc__', '__annotations__')
Thomas Wouters73e5a5b2006-06-08 15:35:45 +000027WRAPPER_UPDATES = ('__dict__',)
28def update_wrapper(wrapper,
29 wrapped,
30 assigned = WRAPPER_ASSIGNMENTS,
31 updated = WRAPPER_UPDATES):
32 """Update a wrapper function to look like the wrapped function
33
34 wrapper is the function to be updated
35 wrapped is the original function
36 assigned is a tuple naming the attributes assigned directly
37 from the wrapped function to the wrapper function (defaults to
38 functools.WRAPPER_ASSIGNMENTS)
Thomas Wouters89f507f2006-12-13 04:49:30 +000039 updated is a tuple naming the attributes of the wrapper that
Thomas Wouters73e5a5b2006-06-08 15:35:45 +000040 are updated with the corresponding attribute from the wrapped
41 function (defaults to functools.WRAPPER_UPDATES)
42 """
43 for attr in assigned:
Antoine Pitrou560f7642010-08-04 18:28:02 +000044 if hasattr(wrapped, attr):
45 setattr(wrapper, attr, getattr(wrapped, attr))
Thomas Wouters73e5a5b2006-06-08 15:35:45 +000046 for attr in updated:
Thomas Wouters89f507f2006-12-13 04:49:30 +000047 getattr(wrapper, attr).update(getattr(wrapped, attr, {}))
Thomas Wouters73e5a5b2006-06-08 15:35:45 +000048 # Return the wrapper so this can be used as a decorator via partial()
49 return wrapper
50
51def wraps(wrapped,
52 assigned = WRAPPER_ASSIGNMENTS,
53 updated = WRAPPER_UPDATES):
54 """Decorator factory to apply update_wrapper() to a wrapper function
55
56 Returns a decorator that invokes update_wrapper() with the decorated
57 function as the wrapper argument and the arguments to wraps() as the
58 remaining arguments. Default arguments are as for update_wrapper().
59 This is a convenience function to simplify applying partial() to
60 update_wrapper().
61 """
62 return partial(update_wrapper, wrapped=wrapped,
63 assigned=assigned, updated=updated)
Raymond Hettingerc50846a2010-04-05 18:56:31 +000064
65def total_ordering(cls):
Georg Brandle5a26732010-05-19 21:06:36 +000066 """Class decorator that fills in missing ordering methods"""
Raymond Hettingerc50846a2010-04-05 18:56:31 +000067 convert = {
68 '__lt__': [('__gt__', lambda self, other: other < self),
69 ('__le__', lambda self, other: not other < self),
70 ('__ge__', lambda self, other: not self < other)],
71 '__le__': [('__ge__', lambda self, other: other <= self),
72 ('__lt__', lambda self, other: not other <= self),
73 ('__gt__', lambda self, other: not self <= other)],
74 '__gt__': [('__lt__', lambda self, other: other > self),
75 ('__ge__', lambda self, other: not other > self),
76 ('__le__', lambda self, other: not self > other)],
77 '__ge__': [('__le__', lambda self, other: other >= self),
78 ('__gt__', lambda self, other: not other >= self),
79 ('__lt__', lambda self, other: not self >= other)]
80 }
81 roots = set(dir(cls)) & set(convert)
Raymond Hettinger56de7e22010-04-10 16:59:03 +000082 if not roots:
83 raise ValueError('must define at least one ordering operation: < > <= >=')
84 root = max(roots) # prefer __lt__ to __le__ to __gt__ to __ge__
Raymond Hettingerc50846a2010-04-05 18:56:31 +000085 for opname, opfunc in convert[root]:
86 if opname not in roots:
87 opfunc.__name__ = opname
88 opfunc.__doc__ = getattr(int, opname).__doc__
89 setattr(cls, opname, opfunc)
90 return cls
91
92def cmp_to_key(mycmp):
Georg Brandle5a26732010-05-19 21:06:36 +000093 """Convert a cmp= function into a key= function"""
Raymond Hettingerc50846a2010-04-05 18:56:31 +000094 class K(object):
95 def __init__(self, obj, *args):
96 self.obj = obj
97 def __lt__(self, other):
98 return mycmp(self.obj, other.obj) < 0
99 def __gt__(self, other):
100 return mycmp(self.obj, other.obj) > 0
101 def __eq__(self, other):
102 return mycmp(self.obj, other.obj) == 0
103 def __le__(self, other):
104 return mycmp(self.obj, other.obj) <= 0
105 def __ge__(self, other):
106 return mycmp(self.obj, other.obj) >= 0
107 def __ne__(self, other):
108 return mycmp(self.obj, other.obj) != 0
109 def __hash__(self):
110 raise TypeError('hash not implemented')
111 return K
Georg Brandl2e7346a2010-07-31 18:09:23 +0000112
Georg Brandl2e7346a2010-07-31 18:09:23 +0000113def lru_cache(maxsize=100):
Benjamin Peterson1f594ad2010-08-08 13:17:07 +0000114 """Least-recently-used cache decorator.
Georg Brandl2e7346a2010-07-31 18:09:23 +0000115
116 Arguments to the cached function must be hashable.
117 Cache performance statistics stored in f.hits and f.misses.
118 Clear the cache using f.clear().
119 http://en.wikipedia.org/wiki/Cache_algorithms#Least_Recently_Used
120
Benjamin Peterson1f594ad2010-08-08 13:17:07 +0000121 """
Raymond Hettinger5202fad2010-08-14 22:29:52 +0000122 def decorating_function(user_function, tuple=tuple, sorted=sorted,
123 len=len, KeyError=KeyError):
Georg Brandl2e7346a2010-07-31 18:09:23 +0000124 cache = OrderedDict() # ordered least recent to most recent
125 kwd_mark = object() # separate positional and keyword args
Raymond Hettingercbe88132010-08-14 22:22:10 +0000126 lock = Lock()
Georg Brandl2e7346a2010-07-31 18:09:23 +0000127
128 @wraps(user_function)
129 def wrapper(*args, **kwds):
130 key = args
131 if kwds:
132 key += (kwd_mark,) + tuple(sorted(kwds.items()))
133 try:
Raymond Hettingercbe88132010-08-14 22:22:10 +0000134 with lock:
135 result = cache[key]
136 del cache[key]
137 cache[key] = result # record recent use of this key
138 wrapper.hits += 1
Georg Brandl2e7346a2010-07-31 18:09:23 +0000139 except KeyError:
140 result = user_function(*args, **kwds)
Raymond Hettingercbe88132010-08-14 22:22:10 +0000141 with lock:
142 cache[key] = result # record recent use of this key
143 wrapper.misses += 1
144 if len(cache) > maxsize:
145 cache.popitem(0) # purge least recently used cache entry
Georg Brandl2e7346a2010-07-31 18:09:23 +0000146 return result
147
148 def clear():
Benjamin Peterson1f594ad2010-08-08 13:17:07 +0000149 """Clear the cache and cache statistics"""
Raymond Hettingercbe88132010-08-14 22:22:10 +0000150 with lock:
151 cache.clear()
152 wrapper.hits = wrapper.misses = 0
Georg Brandl2e7346a2010-07-31 18:09:23 +0000153
154 wrapper.hits = wrapper.misses = 0
155 wrapper.clear = clear
156 return wrapper
157 return decorating_function