decorator.py 17 KB

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