-
Notifications
You must be signed in to change notification settings - Fork 61
/
utils.py
113 lines (86 loc) · 2.8 KB
/
utils.py
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
import sys
from functools import wraps
from textwrap import dedent
try:
from inspect import getfullargspec as getargspec
except ImportError:
from inspect import getargspec
from . import singleton
from .compat import ClassType
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, after):
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, after=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(getargspec(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, objs):
for obj in objs:
try:
obj.__doc__ = dedent(obj.__doc__) + docstring
except (AttributeError, TypeError): # python 2 or pypy 4.0.1 :-(
pass
def indent(text, indent_size = 2):
indented = []
for do_indent, line in enumerate(text.splitlines(True)):
if do_indent:
line = ' '*indent_size + line
indented.append(line)
return ''.join(indented)