decorator.py 17 KB

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