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