import sys from functools import wraps from textwrap import dedent from inspect import getfullargspec from typing import Callable, Sequence, Any from . import singleton DEFAULT = singleton('DEFAULT') defaults = [DEFAULT] try: from .mock import DEFAULT except ImportError: # pragma: no cover pass else: defaults.append(DEFAULT) def generator(*args): """ A utility function for creating a generator that will yield the supplied arguments. """ for i in args: yield i class Wrapping: attribute_name = None new = DEFAULT def __init__(self, before: Callable[[], None], after: Callable[[], None]): self.before, self.after = before, after def __enter__(self): return self.before() def __exit__(self, exc_type=None, exc_val=None, exc_tb=None): if self.after is not None: self.after() def wrap(before: Callable[[], Any], after: Callable[[], Any] = None): """ A decorator that causes the supplied callables to be called before or after the wrapped callable, as appropriate. """ wrapping = Wrapping(before, after) def wrapper(func): if hasattr(func, 'patchings'): func.patchings.append(wrapping) return func @wraps(func) def patched(*args, **keywargs): extra_args = [] entered_patchers = [] to_add = len(getfullargspec(func).args[len(args):]) added = 0 exc_info = (None, None, None) try: for patching in patched.patchings: arg = patching.__enter__() entered_patchers.append(patching) if patching.attribute_name is not None: keywargs.update(arg) elif patching.new in defaults and added < to_add: extra_args.append(arg) added += 1 args += tuple(extra_args) return func(*args, **keywargs) except: # Pass the exception to __exit__ exc_info = sys.exc_info() # re-raise the exception raise finally: for patching in reversed(entered_patchers): patching.__exit__(*exc_info) patched.patchings = [wrapping] return patched return wrapper def extend_docstring(docstring: str, objs: Sequence): for obj in objs: obj.__doc__ = dedent(obj.__doc__) + docstring def indent(text: str, indent_size: int = 2): indented = [] for do_indent, line in enumerate(text.splitlines(True)): if do_indent: line = ' '*indent_size + line indented.append(line) return ''.join(indented)