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

1import collections.abc 

2import contextlib 

3import contextvars 

4import functools 

5import inspect 

6from collections.abc import Callable 

7from contextvars import ContextVar 

8from typing import ParamSpec, TypeVar, NoReturn, Any 

9 

10from debputy.exceptions import ( 

11 UnhandledOrUnexpectedErrorFromPluginError, 

12 DebputyRuntimeError, 

13) 

14from debputy.packages import SourcePackage, BinaryPackage 

15from debputy.util import _trace_log, _is_trace_log_enabled 

16 

17_current_debputy_plugin_cxt_var: ContextVar[str | None] = ContextVar( 

18 "current_debputy_plugin", 

19 default=None, 

20) 

21_current_debputy_parsing_context: ContextVar[ 

22 tuple[ 

23 dict[tuple[SourcePackage | BinaryPackage, type[Any]], Any], 

24 SourcePackage | BinaryPackage, 

25 ] 

26 | None 

27] = ContextVar( 

28 "current_debputy_parsing_context", 

29 default=None, 

30) 

31 

32P = ParamSpec("P") 

33R = TypeVar("R") 

34 

35 

36def register_manifest_type_value_in_context( 

37 value_type: type[Any], 

38 value: Any, 

39) -> None: 

40 context_vars = _current_debputy_parsing_context.get() 

41 if context_vars is None: 

42 raise AssertionError( 

43 "register_manifest_type_value_in_context() was called, but no context was set." 

44 ) 

45 value_table, context_pkg = context_vars 

46 if (context_pkg, value_type) in value_table: 

47 raise AssertionError( 

48 f"The type {value_type!r} was already registered for {context_pkg}, which the plugin API should have prevented" 

49 ) 

50 value_table[(context_pkg, value_type)] = value 

51 

52 

53def begin_parsing_context( 

54 value_table: dict[tuple[SourcePackage | BinaryPackage, type[Any]], Any], 

55 context_pkg: SourcePackage, 

56 func: Callable[P, R], 

57 *args: P.args, 

58 **kwargs: P.kwargs, 

59) -> R: 

60 context = contextvars.copy_context() 

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

62 context.run(_current_debputy_parsing_context.set, (value_table, context_pkg)) 

63 assert context.get(_current_debputy_parsing_context) == (value_table, context_pkg) 

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

65 

66 

67@contextlib.contextmanager 

68def with_binary_pkg_parsing_context( 

69 context_pkg: BinaryPackage, 

70) -> collections.abc.Iterator[None]: 

71 context_vars = _current_debputy_parsing_context.get() 

72 if context_vars is None: 

73 raise AssertionError( 

74 "with_binary_pkg_parsing_context() was called, but no context was set." 

75 ) 

76 value_table, _ = context_vars 

77 token = _current_debputy_parsing_context.set((value_table, context_pkg)) 

78 try: 

79 yield 

80 finally: 

81 _current_debputy_parsing_context.reset(token) 

82 

83 

84def current_debputy_plugin_if_present() -> str | None: 

85 return _current_debputy_plugin_cxt_var.get() 

86 

87 

88def current_debputy_plugin_required() -> str: 

89 v = current_debputy_plugin_if_present() 

90 if v is None: 

91 raise AssertionError( 

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

93 ) 

94 return v 

95 

96 

97def wrap_plugin_code( 

98 plugin_name: str, 

99 func: Callable[P, R], 

100 *, 

101 non_debputy_exception_handling: bool | Callable[[Exception], NoReturn] = True, 

102) -> Callable[P, R]: 

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

104 

105 runner = run_in_context_of_plugin 

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

107 runner = run_in_context_of_plugin_wrap_errors 

108 

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

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

111 

112 functools.update_wrapper(_plugin_wrapper, func) 

113 return _plugin_wrapper 

114 

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

116 try: 

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

118 except DebputyRuntimeError: 

119 raise 

120 except Exception as e: 

121 non_debputy_exception_handling(e) 

122 

123 functools.update_wrapper(_wrapper, func) 

124 return _wrapper 

125 

126 

127def run_in_context_of_plugin( 

128 plugin: str, 

129 func: Callable[P, R], 

130 *args: P.args, 

131 **kwargs: P.kwargs, 

132) -> R: 

133 context = contextvars.copy_context() 

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

135 call_stack = inspect.stack() 

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

137 for frame in call_stack: 

138 if frame.filename != __file__: 

139 try: 

140 fname = frame.frame.f_code.co_qualname 

141 except AttributeError: 

142 fname = None 

143 if fname is None: 

144 fname = frame.function 

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

146 break 

147 # Do not keep the reference longer than necessary 

148 del call_stack 

149 _trace_log( 

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

151 ) 

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

153 context.run(_current_debputy_plugin_cxt_var.set, plugin) 

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

155 

156 

157def run_in_context_of_plugin_wrap_errors( 

158 plugin: str, 

159 func: Callable[P, R], 

160 *args: P.args, 

161 **kwargs: P.kwargs, 

162) -> R: 

163 try: 

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

165 except DebputyRuntimeError: 

166 raise 

167 except Exception as e: 

168 if plugin != "debputy": 

169 raise UnhandledOrUnexpectedErrorFromPluginError( 

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

171 ) from e 

172 else: 

173 raise AssertionError( 

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

175 ) from e