1import sys 2from functools import wraps 3from textwrap import dedent 4 5try: 6 from inspect import getfullargspec as getargspec 7except ImportError: 8 from inspect import getargspec 9 10from . import singleton 11from .compat import ClassType 12 13DEFAULT = singleton('DEFAULT') 14defaults = [DEFAULT] 15 16 17try: 18 from .mock import DEFAULT 19except ImportError: # pragma: no cover 20 pass 21else: 22 defaults.append(DEFAULT) 23 24 25def generator(*args): 26 """ 27 A utility function for creating a generator that will yield the 28 supplied arguments. 29 """ 30 for i in args: 31 yield i 32 33 34class Wrapping: 35 36 attribute_name = None 37 new = DEFAULT 38 39 def __init__(self, before, after): 40 self.before, self.after = before, after 41 42 def __enter__(self): 43 return self.before() 44 45 def __exit__(self, exc_type=None, exc_val=None, exc_tb=None): 46 if self.after is not None: 47 self.after() 48 49 50def wrap(before, after=None): 51 """ 52 A decorator that causes the supplied callables to be called before 53 or after the wrapped callable, as appropriate. 54 """ 55 56 wrapping = Wrapping(before, after) 57 58 def wrapper(func): 59 if hasattr(func, 'patchings'): 60 func.patchings.append(wrapping) 61 return func 62 63 @wraps(func) 64 def patched(*args, **keywargs): 65 extra_args = [] 66 entered_patchers = [] 67 68 to_add = len(getargspec(func).args[len(args):]) 69 added = 0 70 71 exc_info = (None, None, None) 72 try: 73 for patching in patched.patchings: 74 arg = patching.__enter__() 75 entered_patchers.append(patching) 76 if patching.attribute_name is not None: 77 keywargs.update(arg) 78 elif patching.new in defaults and added < to_add: 79 extra_args.append(arg) 80 added += 1 81 82 args += tuple(extra_args) 83 return func(*args, **keywargs) 84 except: 85 # Pass the exception to __exit__ 86 exc_info = sys.exc_info() 87 # re-raise the exception 88 raise 89 finally: 90 for patching in reversed(entered_patchers): 91 patching.__exit__(*exc_info) 92 93 patched.patchings = [wrapping] 94 return patched 95 96 return wrapper 97 98 99def extend_docstring(docstring, objs): 100 for obj in objs: 101 try: 102 obj.__doc__ = dedent(obj.__doc__) + docstring 103 except (AttributeError, TypeError): # python 2 or pypy 4.0.1 :-( 104 pass 105 106 107def indent(text, indent_size = 2): 108 indented = [] 109 for do_indent, line in enumerate(text.splitlines(True)): 110 if do_indent: 111 line = ' '*indent_size + line 112 indented.append(line) 113 return ''.join(indented) 114