decorator.py 15 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365366367368369370371372373374375376377378379380381382383384385386387388389390391392393394395396397398399
  1. # ######################### LICENSE ############################ #
  2. # Copyright (c) 2005-2015, Michele Simionato
  3. # All rights reserved.
  4. # Redistribution and use in source and binary forms, with or without
  5. # modification, are permitted provided that the following conditions are
  6. # met:
  7. # Redistributions of source code must retain the above copyright
  8. # notice, this list of conditions and the following disclaimer.
  9. # Redistributions in bytecode form must reproduce the above copyright
  10. # notice, this list of conditions and the following disclaimer in
  11. # the documentation and/or other materials provided with the
  12. # distribution.
  13. # THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS
  14. # "AS IS" AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT
  15. # LIMITED TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR
  16. # A PARTICULAR PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT
  17. # HOLDERS OR CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT,
  18. # INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING,
  19. # BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS
  20. # OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND
  21. # ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR
  22. # TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE
  23. # USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH
  24. # DAMAGE.
  25. """
  26. Decorator module, see https://pypi.python.org/pypi/decorator
  27. for the documentation.
  28. """
  29. import re
  30. import sys
  31. import inspect
  32. import operator
  33. import itertools
  34. import collections
  35. from inspect import getfullargspec
  36. __version__ = '4.0.5'
  37. def get_init(cls):
  38. return cls.__init__
  39. # getargspec has been deprecated in Python 3.5
  40. ArgSpec = collections.namedtuple(
  41. 'ArgSpec', 'args varargs varkw defaults')
  42. def getargspec(f):
  43. """A replacement for inspect.getargspec"""
  44. spec = getfullargspec(f)
  45. return ArgSpec(spec.args, spec.varargs, spec.varkw, spec.defaults)
  46. DEF = re.compile(r'\s*def\s*([_\w][_\w\d]*)\s*\(')
  47. # basic functionality
  48. class FunctionMaker:
  49. """
  50. An object with the ability to create functions with a given signature.
  51. It has attributes name, doc, module, signature, defaults, dict, and
  52. methods update and make.
  53. """
  54. # Atomic get-and-increment provided by the GIL
  55. _compile_count = itertools.count()
  56. def __init__(self, func=None, name=None, signature=None,
  57. defaults=None, doc=None, module=None, funcdict=None):
  58. self.shortsignature = signature
  59. if func:
  60. # func can be a class or a callable, but not an instance method
  61. self.name = func.__name__
  62. if self.name == '<lambda>': # small hack for lambda functions
  63. self.name = '_lambda_'
  64. self.doc = func.__doc__
  65. self.module = func.__module__
  66. if inspect.isfunction(func):
  67. argspec = getfullargspec(func)
  68. self.annotations = getattr(func, '__annotations__', {})
  69. for a in ('args', 'varargs', 'varkw', 'defaults', 'kwonlyargs',
  70. 'kwonlydefaults'):
  71. setattr(self, a, getattr(argspec, a))
  72. for i, arg in enumerate(self.args):
  73. setattr(self, 'arg%d' % i, arg)
  74. allargs = list(self.args)
  75. allshortargs = list(self.args)
  76. if self.varargs:
  77. allargs.append('*' + self.varargs)
  78. allshortargs.append('*' + self.varargs)
  79. elif self.kwonlyargs:
  80. allargs.append('*') # single star syntax
  81. for a in self.kwonlyargs:
  82. allargs.append('%s=None' % a)
  83. allshortargs.append('%s=%s' % (a, a))
  84. if self.varkw:
  85. allargs.append('**' + self.varkw)
  86. allshortargs.append('**' + self.varkw)
  87. self.signature = ', '.join(allargs)
  88. self.shortsignature = ', '.join(allshortargs)
  89. self.dict = func.__dict__.copy()
  90. # func=None happens when decorating a caller
  91. if name:
  92. self.name = name
  93. if signature is not None:
  94. self.signature = signature
  95. if defaults:
  96. self.defaults = defaults
  97. if doc:
  98. self.doc = doc
  99. if module:
  100. self.module = module
  101. if funcdict:
  102. self.dict = funcdict
  103. # check existence required attributes
  104. assert hasattr(self, 'name')
  105. if not hasattr(self, 'signature'):
  106. raise TypeError('You are decorating a non-function: %s' % func)
  107. def update(self, func, **kw):
  108. "Update the signature of func with the data in self"
  109. func.__name__ = self.name
  110. func.__doc__ = getattr(self, 'doc', None)
  111. func.__dict__ = getattr(self, 'dict', {})
  112. func.__defaults__ = getattr(self, 'defaults', ())
  113. func.__kwdefaults__ = getattr(self, 'kwonlydefaults', None)
  114. func.__annotations__ = getattr(self, 'annotations', None)
  115. try:
  116. frame = sys._getframe(3)
  117. except AttributeError: # for IronPython and similar implementations
  118. callermodule = '?'
  119. else:
  120. callermodule = frame.f_globals.get('__name__', '?')
  121. func.__module__ = getattr(self, 'module', callermodule)
  122. func.__dict__.update(kw)
  123. def make(self, src_templ, evaldict=None, addsource=False, **attrs):
  124. "Make a new function from a given template and update the signature"
  125. src = src_templ % vars(self) # expand name and signature
  126. evaldict = evaldict or {}
  127. mo = DEF.match(src)
  128. if mo is None:
  129. raise SyntaxError('not a valid function template\n%s' % src)
  130. name = mo.group(1) # extract the function name
  131. names = set([name] + [arg.strip(' *') for arg in
  132. self.shortsignature.split(',')])
  133. for n in names:
  134. if n in ('_func_', '_call_'):
  135. raise NameError('%s is overridden in\n%s' % (n, src))
  136. if not src.endswith('\n'): # add a newline just for safety
  137. src += '\n' # this is needed in old versions of Python
  138. # Ensure each generated function has a unique filename for profilers
  139. # (such as cProfile) that depend on the tuple of (<filename>,
  140. # <definition line>, <function name>) being unique.
  141. filename = '<decorator-gen-%d>' % (next(self._compile_count),)
  142. try:
  143. code = compile(src, filename, 'single')
  144. exec(code, evaldict)
  145. except: # noqa: E722
  146. print('Error in generated code:', file=sys.stderr)
  147. print(src, file=sys.stderr)
  148. raise
  149. func = evaldict[name]
  150. if addsource:
  151. attrs['__source__'] = src
  152. self.update(func, **attrs)
  153. return func
  154. @classmethod
  155. def create(cls, obj, body, evaldict, defaults=None,
  156. doc=None, module=None, addsource=True, **attrs):
  157. """
  158. Create a function from the strings name, signature, and body.
  159. evaldict is the evaluation dictionary. If addsource is true, an
  160. attribute __source__ is added to the result. The attributes attrs
  161. are added, if any.
  162. """
  163. if isinstance(obj, str): # "name(signature)"
  164. name, rest = obj.strip().split('(', 1)
  165. signature = rest[:-1] # strip a right parens
  166. func = None
  167. else: # a function
  168. name = None
  169. signature = None
  170. func = obj
  171. self = cls(func, name, signature, defaults, doc, module)
  172. ibody = '\n'.join(' ' + line for line in body.splitlines())
  173. return self.make('def %(name)s(%(signature)s):\n' + ibody,
  174. evaldict, addsource, **attrs)
  175. def decorate(func, caller):
  176. """
  177. decorate(func, caller) decorates a function using a caller.
  178. """
  179. evaldict = func.__globals__.copy()
  180. evaldict['_call_'] = caller
  181. evaldict['_func_'] = func
  182. fun = FunctionMaker.create(
  183. func, "return _call_(_func_, %(shortsignature)s)",
  184. evaldict, __wrapped__=func)
  185. if hasattr(func, '__qualname__'):
  186. fun.__qualname__ = func.__qualname__
  187. return fun
  188. def decorator(caller, _func=None):
  189. """decorator(caller) converts a caller function into a decorator"""
  190. if _func is not None: # return a decorated function
  191. # this is obsolete behavior; you should use decorate instead
  192. return decorate(_func, caller)
  193. # else return a decorator function
  194. if inspect.isclass(caller):
  195. name = caller.__name__.lower()
  196. callerfunc = get_init(caller)
  197. doc = 'decorator(%s) converts functions/generators into ' \
  198. 'factories of %s objects' % (caller.__name__, caller.__name__)
  199. elif inspect.isfunction(caller):
  200. if caller.__name__ == '<lambda>':
  201. name = '_lambda_'
  202. else:
  203. name = caller.__name__
  204. callerfunc = caller
  205. doc = caller.__doc__
  206. else: # assume caller is an object with a __call__ method
  207. name = caller.__class__.__name__.lower()
  208. callerfunc = caller.__call__.__func__
  209. doc = caller.__call__.__doc__
  210. evaldict = callerfunc.__globals__.copy()
  211. evaldict['_call_'] = caller
  212. evaldict['_decorate_'] = decorate
  213. return FunctionMaker.create(
  214. '%s(func)' % name, 'return _decorate_(func, _call_)',
  215. evaldict, doc=doc, module=caller.__module__,
  216. __wrapped__=caller)
  217. # ####################### contextmanager ####################### #
  218. try: # Python >= 3.2
  219. from contextlib import _GeneratorContextManager
  220. except ImportError: # Python >= 2.5
  221. from contextlib import GeneratorContextManager as _GeneratorContextManager
  222. class ContextManager(_GeneratorContextManager):
  223. def __call__(self, func):
  224. """Context manager decorator"""
  225. return FunctionMaker.create(
  226. func, "with _self_: return _func_(%(shortsignature)s)",
  227. dict(_self_=self, _func_=func), __wrapped__=func)
  228. init = getfullargspec(_GeneratorContextManager.__init__)
  229. n_args = len(init.args)
  230. if n_args == 2 and not init.varargs: # (self, genobj) Python 2.7
  231. def __init__(self, g, *a, **k):
  232. return _GeneratorContextManager.__init__(self, g(*a, **k))
  233. ContextManager.__init__ = __init__
  234. elif n_args == 2 and init.varargs: # (self, gen, *a, **k) Python 3.4
  235. pass
  236. elif n_args == 4: # (self, gen, args, kwds) Python 3.5
  237. def __init__(self, g, *a, **k):
  238. return _GeneratorContextManager.__init__(self, g, a, k)
  239. ContextManager.__init__ = __init__
  240. contextmanager = decorator(ContextManager)
  241. # ############################ dispatch_on ############################ #
  242. def append(a, vancestors):
  243. """
  244. Append ``a`` to the list of the virtual ancestors, unless it is already
  245. included.
  246. """
  247. add = True
  248. for j, va in enumerate(vancestors):
  249. if issubclass(va, a):
  250. add = False
  251. break
  252. if issubclass(a, va):
  253. vancestors[j] = a
  254. add = False
  255. if add:
  256. vancestors.append(a)
  257. # inspired from simplegeneric by P.J. Eby and functools.singledispatch
  258. def dispatch_on(*dispatch_args):
  259. """
  260. Factory of decorators turning a function into a generic function
  261. dispatching on the given arguments.
  262. """
  263. assert dispatch_args, 'No dispatch args passed'
  264. dispatch_str = '(%s,)' % ', '.join(dispatch_args)
  265. def check(arguments, wrong=operator.ne, msg=''):
  266. """Make sure one passes the expected number of arguments"""
  267. if wrong(len(arguments), len(dispatch_args)):
  268. raise TypeError('Expected %d arguments, got %d%s' %
  269. (len(dispatch_args), len(arguments), msg))
  270. def gen_func_dec(func):
  271. """Decorator turning a function into a generic function"""
  272. # first check the dispatch arguments
  273. argset = set(getfullargspec(func).args)
  274. if not set(dispatch_args) <= argset:
  275. raise NameError('Unknown dispatch arguments %s' % dispatch_str)
  276. typemap = {}
  277. def vancestors(*types):
  278. """
  279. Get a list of sets of virtual ancestors for the given types
  280. """
  281. check(types)
  282. ras = [[] for _ in range(len(dispatch_args))]
  283. for types_ in typemap:
  284. for t, type_, ra in zip(types, types_, ras):
  285. if issubclass(t, type_) and type_ not in t.__mro__:
  286. append(type_, ra)
  287. return [set(ra) for ra in ras]
  288. def ancestors(*types):
  289. """
  290. Get a list of virtual MROs, one for each type
  291. """
  292. check(types)
  293. lists = []
  294. for t, vas in zip(types, vancestors(*types)):
  295. n_vas = len(vas)
  296. if n_vas > 1:
  297. raise RuntimeError(
  298. 'Ambiguous dispatch for %s: %s' % (t, vas))
  299. elif n_vas == 1:
  300. va, = vas
  301. mro = type('t', (t, va), {}).__mro__[1:]
  302. else:
  303. mro = t.__mro__
  304. lists.append(mro[:-1]) # discard t and object
  305. return lists
  306. def register(*types):
  307. """
  308. Decorator to register an implementation for the given types
  309. """
  310. check(types)
  311. def dec(f):
  312. check(getfullargspec(f).args, operator.lt, ' in ' + f.__name__)
  313. typemap[types] = f
  314. return f
  315. return dec
  316. def dispatch_info(*types):
  317. """
  318. An utility to introspect the dispatch algorithm
  319. """
  320. check(types)
  321. lst = [tuple(a.__name__ for a in anc)
  322. for anc in itertools.product(*ancestors(*types))]
  323. return lst
  324. def _dispatch(dispatch_args, *args, **kw):
  325. types = tuple(type(arg) for arg in dispatch_args)
  326. try: # fast path
  327. f = typemap[types]
  328. except KeyError:
  329. pass
  330. else:
  331. return f(*args, **kw)
  332. combinations = itertools.product(*ancestors(*types))
  333. next(combinations) # the first one has been already tried
  334. for types_ in combinations:
  335. f = typemap.get(types_)
  336. if f is not None:
  337. return f(*args, **kw)
  338. # else call the default implementation
  339. return func(*args, **kw)
  340. return FunctionMaker.create(
  341. func, 'return _f_(%s, %%(shortsignature)s)' % dispatch_str,
  342. dict(_f_=_dispatch), register=register, default=func,
  343. typemap=typemap, vancestors=vancestors, ancestors=ancestors,
  344. dispatch_info=dispatch_info, __wrapped__=func)
  345. gen_func_dec.__name__ = 'dispatch_on' + dispatch_str
  346. return gen_func_dec