|
1
|
|
|
import functools |
|
2
|
|
|
|
|
3
|
|
|
|
|
4
|
|
|
class ContextMethodDecorator(): |
|
5
|
|
|
"""A helper ContextManager decorating a method with a custom function.""" |
|
6
|
|
|
|
|
7
|
|
|
def __init__(self, classx, method_name, decorator_func): |
|
8
|
|
|
""" |
|
9
|
|
|
Create a new context manager decorating a function within its scope. |
|
10
|
|
|
|
|
11
|
|
|
This is a helper Context Manager that decorates a method of a class |
|
12
|
|
|
with a custom function. |
|
13
|
|
|
The decoration is only valid within the scope. |
|
14
|
|
|
:param classx: A class (object) |
|
15
|
|
|
:param method_name A string name of the method to be decorated |
|
16
|
|
|
:param decorator_func: The decorator function is responsible |
|
17
|
|
|
for calling the original method. |
|
18
|
|
|
The signature should be: func(instance, original_method, |
|
19
|
|
|
original_args, original_kwargs) |
|
20
|
|
|
when called, instance refers to an instance of classx and the |
|
21
|
|
|
original_method refers to the original method object which can be |
|
22
|
|
|
called. |
|
23
|
|
|
args and kwargs are arguments passed to the method |
|
24
|
|
|
|
|
25
|
|
|
""" |
|
26
|
|
|
self.method_name = method_name |
|
27
|
|
|
self.decorator_func = decorator_func |
|
28
|
|
|
self.classx = classx |
|
29
|
|
|
self.patched_by_me = False |
|
30
|
|
|
|
|
31
|
|
|
def __enter__(self): |
|
32
|
|
|
|
|
33
|
|
|
self.original_method = getattr(self.classx, self.method_name) |
|
34
|
|
|
if not getattr(self.original_method, |
|
35
|
|
|
"sacred_patched%s" % self.__class__.__name__, False): |
|
36
|
|
|
@functools.wraps(self.original_method) |
|
37
|
|
|
def decorated(instance, *args, **kwargs): |
|
38
|
|
|
return self.decorator_func(instance, |
|
39
|
|
|
self.original_method, args, |
|
40
|
|
|
kwargs) |
|
41
|
|
|
|
|
42
|
|
|
setattr(self.classx, self.method_name, decorated) |
|
43
|
|
|
setattr(decorated, |
|
44
|
|
|
"sacred_patched%s" % self.__class__.__name__, True) |
|
45
|
|
|
self.patched_by_me = True |
|
46
|
|
|
|
|
47
|
|
|
def __exit__(self, type, value, traceback): |
|
48
|
|
|
if self.patched_by_me: |
|
49
|
|
|
# Restore original function |
|
50
|
|
|
setattr(self.classx, self.method_name, self.original_method) |
|
51
|
|
|
|
IDSIA /
sacred
