Coverage for sacred/sacred/stflow/internal.py: 0%

Shortcuts on this page

r m x   toggle line displays

j k   next/prev highlighted chunk

0   (zero) top of page

1   (one) first highlighted chunk

19 statements  

1import functools 

2 

3 

4class 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 hasattr( 

35 self.original_method, "sacred_patched%s" % self.__class__.__name__ 

36 ): 

37 

38 @functools.wraps(self.original_method) 

39 def decorated(instance, *args, **kwargs): 

40 return self.decorator_func(instance, self.original_method, args, kwargs) 

41 

42 setattr(self.classx, self.method_name, decorated) 

43 setattr(decorated, "sacred_patched%s" % self.__class__.__name__, True) 

44 self.patched_by_me = True 

45 

46 def __exit__(self, type, value, traceback): 

47 if self.patched_by_me: 

48 # Restore original function 

49 setattr(self.classx, self.method_name, self.original_method)