decorator.py 16 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365366367368369370371372373374375376377378379380381382383384385386387388389390391392393394395396397398399400401402403404405406407408409410411412413414415416417418419420421422423424425426
  1. # ######################### LICENSE ############################ #
  2. # Copyright (c) 2005-2017, 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 http://pypi.python.org/pypi/decorator
  27. for the documentation.
  28. """
  29. from __future__ import print_function
  30. import re
  31. import sys
  32. import inspect
  33. import operator
  34. import itertools
  35. import collections
  36. __version__ = '4.1.2'
  37. if sys.version >= '3':
  38. from inspect import getfullargspec
  39. def get_init(cls):
  40. return cls.__init__
  41. else:
  42. FullArgSpec = collections.namedtuple(
  43. 'FullArgSpec', 'args varargs varkw defaults '
  44. 'kwonlyargs kwonlydefaults')
  45. def getfullargspec(f):
  46. "A quick and dirty replacement for getfullargspec for Python 2.X"
  47. return FullArgSpec._make(inspect.getargspec(f) + ([], None))
  48. def get_init(cls):
  49. return cls.__init__.__func__
  50. try:
  51. iscoroutinefunction = inspect.iscoroutinefunction
  52. except AttributeError:
  53. # let's assume there are no coroutine functions in old Python
  54. def iscoroutinefunction(f):
  55. return False
  56. # getargspec has been deprecated in Python 3.5
  57. ArgSpec = collections.namedtuple(
  58. 'ArgSpec', 'args varargs varkw defaults')
  59. def getargspec(f):
  60. """A replacement for inspect.getargspec"""
  61. spec = getfullargspec(f)
  62. return ArgSpec(spec.args, spec.varargs, spec.varkw, spec.defaults)
  63. DEF = re.compile(r'\s*def\s*([_\w][_\w\d]*)\s*\(')
  64. # basic functionality
  65. class FunctionMaker(object):
  66. """
  67. An object with the ability to create functions with a given signature.
  68. It has attributes name, doc, module, signature, defaults, dict and
  69. methods update and make.
  70. """
  71. # Atomic get-and-increment provided by the GIL
  72. _compile_count = itertools.count()
  73. # make pylint happy
  74. args = varargs = varkw = defaults = kwonlyargs = kwonlydefaults = ()
  75. def __init__(self, func=None, name=None, signature=None,
  76. defaults=None, doc=None, module=None, funcdict=None):
  77. self.shortsignature = signature
  78. if func:
  79. # func can be a class or a callable, but not an instance method
  80. self.name = func.__name__
  81. if self.name == '<lambda>': # small hack for lambda functions
  82. self.name = '_lambda_'
  83. self.doc = func.__doc__
  84. self.module = func.__module__
  85. if inspect.isfunction(func):
  86. argspec = getfullargspec(func)
  87. self.annotations = getattr(func, '__annotations__', {})
  88. for a in ('args', 'varargs', 'varkw', 'defaults', 'kwonlyargs',
  89. 'kwonlydefaults'):
  90. setattr(self, a, getattr(argspec, a))
  91. for i, arg in enumerate(self.args):
  92. setattr(self, 'arg%d' % i, arg)
  93. if sys.version < '3': # easy way
  94. self.shortsignature = self.signature = (
  95. inspect.formatargspec(
  96. formatvalue=lambda val: "", *argspec[:-2])[1:-1])
  97. else: # Python 3 way
  98. allargs = list(self.args)
  99. allshortargs = list(self.args)
  100. if self.varargs:
  101. allargs.append('*' + self.varargs)
  102. allshortargs.append('*' + self.varargs)
  103. elif self.kwonlyargs:
  104. allargs.append('*') # single star syntax
  105. for a in self.kwonlyargs:
  106. allargs.append('%s=None' % a)
  107. allshortargs.append('%s=%s' % (a, a))
  108. if self.varkw:
  109. allargs.append('**' + self.varkw)
  110. allshortargs.append('**' + self.varkw)
  111. self.signature = ', '.join(allargs)
  112. self.shortsignature = ', '.join(allshortargs)
  113. self.dict = func.__dict__.copy()
  114. # func=None happens when decorating a caller
  115. if name:
  116. self.name = name
  117. if signature is not None:
  118. self.signature = signature
  119. if defaults:
  120. self.defaults = defaults
  121. if doc:
  122. self.doc = doc
  123. if module:
  124. self.module = module
  125. if funcdict:
  126. self.dict = funcdict
  127. # check existence required attributes
  128. assert hasattr(self, 'name')
  129. if not hasattr(self, 'signature'):
  130. raise TypeError('You are decorating a non function: %s' % func)
  131. def update(self, func, **kw):
  132. "Update the signature of func with the data in self"
  133. func.__name__ = self.name
  134. func.__doc__ = getattr(self, 'doc', None)
  135. func.__dict__ = getattr(self, 'dict', {})
  136. func.__defaults__ = self.defaults
  137. func.__kwdefaults__ = self.kwonlydefaults or None
  138. func.__annotations__ = getattr(self, 'annotations', None)
  139. try:
  140. frame = sys._getframe(3)
  141. except AttributeError: # for IronPython and similar implementations
  142. callermodule = '?'
  143. else:
  144. callermodule = frame.f_globals.get('__name__', '?')
  145. func.__module__ = getattr(self, 'module', callermodule)
  146. func.__dict__.update(kw)
  147. def make(self, src_templ, evaldict=None, addsource=False, **attrs):
  148. "Make a new function from a given template and update the signature"
  149. src = src_templ % vars(self) # expand name and signature
  150. evaldict = evaldict or {}
  151. mo = DEF.search(src)
  152. if mo is None:
  153. raise SyntaxError('not a valid function template\n%s' % src)
  154. name = mo.group(1) # extract the function name
  155. names = set([name] + [arg.strip(' *') for arg in
  156. self.shortsignature.split(',')])
  157. for n in names:
  158. if n in ('_func_', '_call_'):
  159. raise NameError('%s is overridden in\n%s' % (n, src))
  160. if not src.endswith('\n'): # add a newline for old Pythons
  161. src += '\n'
  162. # Ensure each generated function has a unique filename for profilers
  163. # (such as cProfile) that depend on the tuple of (<filename>,
  164. # <definition line>, <function name>) being unique.
  165. filename = '<decorator-gen-%d>' % (next(self._compile_count),)
  166. try:
  167. code = compile(src, filename, 'single')
  168. exec(code, evaldict)
  169. except:
  170. print('Error in generated code:', file=sys.stderr)
  171. print(src, file=sys.stderr)
  172. raise
  173. func = evaldict[name]
  174. if addsource:
  175. attrs['__source__'] = src
  176. self.update(func, **attrs)
  177. return func
  178. @classmethod
  179. def create(cls, obj, body, evaldict, defaults=None,
  180. doc=None, module=None, addsource=True, **attrs):
  181. """
  182. Create a function from the strings name, signature and body.
  183. evaldict is the evaluation dictionary. If addsource is true an
  184. attribute __source__ is added to the result. The attributes attrs
  185. are added, if any.
  186. """
  187. if isinstance(obj, str): # "name(signature)"
  188. name, rest = obj.strip().split('(', 1)
  189. signature = rest[:-1] # strip a right parens
  190. func = None
  191. else: # a function
  192. name = None
  193. signature = None
  194. func = obj
  195. self = cls(func, name, signature, defaults, doc, module)
  196. ibody = '\n'.join(' ' + line for line in body.splitlines())
  197. caller = evaldict.get('_call_') # when called from `decorate`
  198. if caller and iscoroutinefunction(caller):
  199. body = ('async def %(name)s(%(signature)s):\n' + ibody).replace(
  200. 'return', 'return await')
  201. else:
  202. body = 'def %(name)s(%(signature)s):\n' + ibody
  203. return self.make(body, evaldict, addsource, **attrs)
  204. def decorate(func, caller):
  205. """
  206. decorate(func, caller) decorates a function using a caller.
  207. """
  208. evaldict = dict(_call_=caller, _func_=func)
  209. fun = FunctionMaker.create(
  210. func, "return _call_(_func_, %(shortsignature)s)",
  211. evaldict, __wrapped__=func)
  212. if hasattr(func, '__qualname__'):
  213. fun.__qualname__ = func.__qualname__
  214. return fun
  215. def decorator(caller, _func=None):
  216. """decorator(caller) converts a caller function into a decorator"""
  217. if _func is not None: # return a decorated function
  218. # this is obsolete behavior; you should use decorate instead
  219. return decorate(_func, caller)
  220. # else return a decorator function
  221. if inspect.isclass(caller):
  222. name = caller.__name__.lower()
  223. doc = 'decorator(%s) converts functions/generators into ' \
  224. 'factories of %s objects' % (caller.__name__, caller.__name__)
  225. elif inspect.isfunction(caller):
  226. if caller.__name__ == '<lambda>':
  227. name = '_lambda_'
  228. else:
  229. name = caller.__name__
  230. doc = caller.__doc__
  231. else: # assume caller is an object with a __call__ method
  232. name = caller.__class__.__name__.lower()
  233. doc = caller.__call__.__doc__
  234. evaldict = dict(_call=caller, _decorate_=decorate)
  235. return FunctionMaker.create(
  236. '%s(func)' % name, 'return _decorate_(func, _call)',
  237. evaldict, doc=doc, module=caller.__module__,
  238. __wrapped__=caller)
  239. # ####################### contextmanager ####################### #
  240. try: # Python >= 3.2
  241. from contextlib import _GeneratorContextManager
  242. except ImportError: # Python >= 2.5
  243. from contextlib import GeneratorContextManager as _GeneratorContextManager
  244. class ContextManager(_GeneratorContextManager):
  245. def __call__(self, func):
  246. """Context manager decorator"""
  247. return FunctionMaker.create(
  248. func, "with _self_: return _func_(%(shortsignature)s)",
  249. dict(_self_=self, _func_=func), __wrapped__=func)
  250. init = getfullargspec(_GeneratorContextManager.__init__)
  251. n_args = len(init.args)
  252. if n_args == 2 and not init.varargs: # (self, genobj) Python 2.7
  253. def __init__(self, g, *a, **k):
  254. return _GeneratorContextManager.__init__(self, g(*a, **k))
  255. ContextManager.__init__ = __init__
  256. elif n_args == 2 and init.varargs: # (self, gen, *a, **k) Python 3.4
  257. pass
  258. elif n_args == 4: # (self, gen, args, kwds) Python 3.5
  259. def __init__(self, g, *a, **k):
  260. return _GeneratorContextManager.__init__(self, g, a, k)
  261. ContextManager.__init__ = __init__
  262. contextmanager = decorator(ContextManager)
  263. # ############################ dispatch_on ############################ #
  264. def append(a, vancestors):
  265. """
  266. Append ``a`` to the list of the virtual ancestors, unless it is already
  267. included.
  268. """
  269. add = True
  270. for j, va in enumerate(vancestors):
  271. if issubclass(va, a):
  272. add = False
  273. break
  274. if issubclass(a, va):
  275. vancestors[j] = a
  276. add = False
  277. if add:
  278. vancestors.append(a)
  279. # inspired from simplegeneric by P.J. Eby and functools.singledispatch
  280. def dispatch_on(*dispatch_args):
  281. """
  282. Factory of decorators turning a function into a generic function
  283. dispatching on the given arguments.
  284. """
  285. assert dispatch_args, 'No dispatch args passed'
  286. dispatch_str = '(%s,)' % ', '.join(dispatch_args)
  287. def check(arguments, wrong=operator.ne, msg=''):
  288. """Make sure one passes the expected number of arguments"""
  289. if wrong(len(arguments), len(dispatch_args)):
  290. raise TypeError('Expected %d arguments, got %d%s' %
  291. (len(dispatch_args), len(arguments), msg))
  292. def gen_func_dec(func):
  293. """Decorator turning a function into a generic function"""
  294. # first check the dispatch arguments
  295. argset = set(getfullargspec(func).args)
  296. if not set(dispatch_args) <= argset:
  297. raise NameError('Unknown dispatch arguments %s' % dispatch_str)
  298. typemap = {}
  299. def vancestors(*types):
  300. """
  301. Get a list of sets of virtual ancestors for the given types
  302. """
  303. check(types)
  304. ras = [[] for _ in range(len(dispatch_args))]
  305. for types_ in typemap:
  306. for t, type_, ra in zip(types, types_, ras):
  307. if issubclass(t, type_) and type_ not in t.mro():
  308. append(type_, ra)
  309. return [set(ra) for ra in ras]
  310. def ancestors(*types):
  311. """
  312. Get a list of virtual MROs, one for each type
  313. """
  314. check(types)
  315. lists = []
  316. for t, vas in zip(types, vancestors(*types)):
  317. n_vas = len(vas)
  318. if n_vas > 1:
  319. raise RuntimeError(
  320. 'Ambiguous dispatch for %s: %s' % (t, vas))
  321. elif n_vas == 1:
  322. va, = vas
  323. mro = type('t', (t, va), {}).mro()[1:]
  324. else:
  325. mro = t.mro()
  326. lists.append(mro[:-1]) # discard t and object
  327. return lists
  328. def register(*types):
  329. """
  330. Decorator to register an implementation for the given types
  331. """
  332. check(types)
  333. def dec(f):
  334. check(getfullargspec(f).args, operator.lt, ' in ' + f.__name__)
  335. typemap[types] = f
  336. return f
  337. return dec
  338. def dispatch_info(*types):
  339. """
  340. An utility to introspect the dispatch algorithm
  341. """
  342. check(types)
  343. lst = []
  344. for anc in itertools.product(*ancestors(*types)):
  345. lst.append(tuple(a.__name__ for a in anc))
  346. return lst
  347. def _dispatch(dispatch_args, *args, **kw):
  348. types = tuple(type(arg) for arg in dispatch_args)
  349. try: # fast path
  350. f = typemap[types]
  351. except KeyError:
  352. pass
  353. else:
  354. return f(*args, **kw)
  355. combinations = itertools.product(*ancestors(*types))
  356. next(combinations) # the first one has been already tried
  357. for types_ in combinations:
  358. f = typemap.get(types_)
  359. if f is not None:
  360. return f(*args, **kw)
  361. # else call the default implementation
  362. return func(*args, **kw)
  363. return FunctionMaker.create(
  364. func, 'return _f_(%s, %%(shortsignature)s)' % dispatch_str,
  365. dict(_f_=_dispatch), register=register, default=func,
  366. typemap=typemap, vancestors=vancestors, ancestors=ancestors,
  367. dispatch_info=dispatch_info, __wrapped__=func)
  368. gen_func_dec.__name__ = 'dispatch_on' + dispatch_str
  369. return gen_func_dec