Coverage for src/debputy/plugin/plugin_state.py: 51%

1import contextvars 

2import functools 

3import inspect 

4from contextvars import ContextVar 

5from typing import Optional, Callable, ParamSpec, TypeVar, NoReturn, Union 

6 

7from debputy.exceptions import ( 

8 UnhandledOrUnexpectedErrorFromPluginError, 

9 DebputyRuntimeError, 

10) 

11from debputy.util import _trace_log, _is_trace_log_enabled 

12 

13_current_debputy_plugin_cxt_var: ContextVar[Optional[str]] = ContextVar( 

14 "current_debputy_plugin", 

15 default=None, 

16) 

17 

18P = ParamSpec("P") 

19R = TypeVar("R") 

20 

21 

22def current_debputy_plugin_if_present() -> Optional[str]: 

23 return _current_debputy_plugin_cxt_var.get() 

24 

25 

26def current_debputy_plugin_required() -> str: 

27 v = current_debputy_plugin_if_present() 

28 if v is None: 

29 raise AssertionError( 

30 "current_debputy_plugin_required() was called, but no plugin was set." 

31 ) 

32 return v 

33 

34 

35def wrap_plugin_code( 

36 plugin_name: str, 

37 func: Callable[P, R], 

38 *, 

39 non_debputy_exception_handling: Union[bool, Callable[[Exception], NoReturn]] = True, 

40) -> Callable[P, R]: 

41 if isinstance(non_debputy_exception_handling, bool): 41 ↛ 53line 41 didn't jump to line 53 because the condition on line 41 was always true

42 

43 runner = run_in_context_of_plugin 

44 if non_debputy_exception_handling: 44 ↛ 47line 44 didn't jump to line 47 because the condition on line 44 was always true

45 runner = run_in_context_of_plugin_wrap_errors 

46 

47 def _plugin_wrapper(*args: P.args, **kwargs: P.kwargs) -> R: 

48 return runner(plugin_name, func, *args, **kwargs) 

49 

50 functools.update_wrapper(_plugin_wrapper, func) 

51 return _plugin_wrapper 

52 

53 def _wrapper(*args: P.args, **kwargs: P.kwargs) -> R: 

54 try: 

55 return run_in_context_of_plugin(plugin_name, func, *args, **kwargs) 

56 except DebputyRuntimeError: 

57 raise 

58 except Exception as e: 

59 non_debputy_exception_handling(e) 

60 

61 functools.update_wrapper(_wrapper, func) 

62 return _wrapper 

63 

64 

65def run_in_context_of_plugin( 

66 plugin: str, 

67 func: Callable[P, R], 

68 *args: P.args, 

69 **kwargs: P.kwargs, 

70) -> R: 

71 context = contextvars.copy_context() 

72 if _is_trace_log_enabled(): 72 ↛ 73line 72 didn't jump to line 73 because the condition on line 72 was never true

73 call_stack = inspect.stack() 

74 caller: str = "[N/A]" 

75 for frame in call_stack: 

76 if frame.filename != __file__: 

77 try: 

78 fname = frame.frame.f_code.co_qualname 

79 except AttributeError: 

80 fname = None 

81 if fname is None: 

82 fname = frame.function 

83 caller = f"{frame.filename}:{frame.lineno} ({fname})" 

84 break 

85 # Do not keep the reference longer than necessary 

86 del call_stack 

87 _trace_log( 

88 f"Switching plugin context to {plugin} at {caller} (from context: {current_debputy_plugin_if_present()})" 

89 ) 

90 # Wish we could just do a regular set without wrapping it in `context.run` 

91 context.run(_current_debputy_plugin_cxt_var.set, plugin) 

92 return context.run(func, *args, **kwargs) 

93 

94 

95def run_in_context_of_plugin_wrap_errors( 

96 plugin: str, 

97 func: Callable[P, R], 

98 *args: P.args, 

99 **kwargs: P.kwargs, 

100) -> R: 

101 try: 

102 return run_in_context_of_plugin(plugin, func, *args, **kwargs) 

103 except DebputyRuntimeError: 103 ↛ 105line 103 didn't jump to line 105

104 raise 

105 except Exception as e: 

106 if plugin != "debputy": 

107 raise UnhandledOrUnexpectedErrorFromPluginError( 

108 f"{func.__qualname__} from the plugin {plugin} raised exception that was not expected here." 

109 ) from e 

110 else: 

111 raise AssertionError( 

112 "Bug in the `debputy` plugin: Unhandled exception." 

113 ) from e