Coverage for /home/runner/work/bijux-cli/bijux-cli/src/bijux_cli/core/context.py: 100%

1# SPDX-License-Identifier: MIT 

2# Copyright © 2025 Bijan Mousavi 

3 

4"""Provides a concrete implementation for request-scoped context management. 

5 

6This module defines the `Context` class, which implements the `ContextProtocol`. 

7It uses Python's `contextvars` to provide a thread-safe and async-safe 

8mechanism for storing and retrieving key-value data associated with a specific 

9command execution or request. This allows state to be passed through the 

10application's call stack without explicit argument passing. 

11""" 

12 

13from __future__ import annotations 

14 

15from collections.abc import Iterator 

16from contextlib import contextmanager 

17from contextvars import ContextVar, Token 

18import os 

19from typing import Any, TypeVar 

20 

21from injector import inject 

22 

23from bijux_cli.contracts import ContextProtocol, ObservabilityProtocol 

24from bijux_cli.core.di import DIContainer 

25 

26T = TypeVar("T") 

27_current_context: ContextVar[dict[str, Any] | None] = ContextVar( 

28 "current_context", default=None 

29) 

30 

31 

32class Context(ContextProtocol): 

33 """Provides thread-safe, request-scoped storage for CLI commands. 

34 

35 This class uses `contextvars` to manage a dictionary of data that is 

36 isolated to the current task or thread. It is intended to be used as 

37 both a synchronous and asynchronous context manager. 

38 

39 Attributes: 

40 _di (DIContainer): The dependency injection container. 

41 _log (ObservabilityProtocol): The logging service. 

42 _data (dict[str, Any]): The dictionary storing the context's data. 

43 _token (Token | None): The token for resetting the `ContextVar`. 

44 """ 

45 

46 @inject 

47 def __init__(self, di: DIContainer) -> None: 

48 """Initializes a new Context instance. 

49 

50 Args: 

51 di (DIContainer): The dependency injection container used to 

52 resolve the logging service. 

53 """ 

54 self._di = di 

55 self._log: ObservabilityProtocol = di.resolve(ObservabilityProtocol) 

56 self._data: dict[str, Any] = {} 

57 self._token: Token[dict[str, Any] | None] | None = None 

58 if os.getenv("VERBOSE_DI") and not os.getenv("BIJUXCLI_TEST_MODE"): 

59 self._log.log("debug", "Context initialized", extra={}) 

60 

61 def set(self, key: str, value: Any) -> None: 

62 """Sets a value in the current context's data. 

63 

64 Args: 

65 key (str): The key for the value. 

66 value (Any): The value to store. 

67 

68 Returns: 

69 None: 

70 """ 

71 self._data[key] = value 

72 if os.getenv("VERBOSE_DI") and not os.getenv("BIJUXCLI_TEST_MODE"): 

73 self._log.log( 

74 "debug", "Context set", extra={"key": key, "value": str(value)} 

75 ) 

76 

77 def get(self, key: str) -> Any: 

78 """Gets a value from the current context's data. 

79 

80 Args: 

81 key (str): The key of the value to retrieve. 

82 

83 Returns: 

84 Any: The value associated with the key. 

85 

86 Raises: 

87 KeyError: If the key is not found in the context. 

88 """ 

89 if key not in self._data: 

90 if os.getenv("VERBOSE_DI") and not os.getenv("BIJUXCLI_TEST_MODE"): 

91 self._log.log("warning", "Context key not found", extra={"key": key}) 

92 raise KeyError(f"Key '{key}' not found in context") 

93 if os.getenv("VERBOSE_DI") and not os.getenv("BIJUXCLI_TEST_MODE"): 

94 self._log.log( 

95 "debug", 

96 "Context get", 

97 extra={"key": key, "value": str(self._data[key])}, 

98 ) 

99 return self._data[key] 

100 

101 def clear(self) -> None: 

102 """Removes all values from the context's data.""" 

103 self._data.clear() 

104 if os.getenv("VERBOSE_DI") and not os.getenv("BIJUXCLI_TEST_MODE"): 

105 self._log.log("debug", "Context cleared", extra={}) 

106 

107 def __enter__(self) -> Context: 

108 """Enters the context as a synchronous manager. 

109 

110 This sets the current `contextvar` to this instance's data dictionary. 

111 

112 Returns: 

113 Context: The context instance itself. 

114 """ 

115 self._token = _current_context.set(self._data) 

116 if os.getenv("VERBOSE_DI") and not os.getenv("BIJUXCLI_TEST_MODE"): 

117 self._log.log("debug", "Context entered", extra={}) 

118 return self 

119 

120 def __exit__(self, _exc_type: Any, _exc_value: Any, traceback: Any) -> None: 

121 """Exits the synchronous context manager. 

122 

123 This resets the `contextvar` to its previous state. 

124 

125 Args: 

126 _exc_type (Any): Exception type, if any (unused). 

127 _exc_value (Any): Exception value, if any (unused). 

128 traceback (Any): Traceback, if any (unused). 

129 

130 Returns: 

131 None: 

132 """ 

133 if self._token: 

134 _current_context.reset(self._token) 

135 self._token = None 

136 if os.getenv("VERBOSE_DI") and not os.getenv("BIJUXCLI_TEST_MODE"): 

137 self._log.log("debug", "Context exited", extra={}) 

138 

139 async def __aenter__(self) -> Context: 

140 """Enters the context as an asynchronous manager. 

141 

142 This sets the current `contextvar` to this instance's data dictionary. 

143 

144 Returns: 

145 Context: The context instance itself. 

146 """ 

147 self._token = _current_context.set(self._data) 

148 if os.getenv("VERBOSE_DI") and not os.getenv("BIJUXCLI_TEST_MODE"): 

149 self._log.log("debug", "Async context entered", extra={}) 

150 return self 

151 

152 async def __aexit__(self, _exc_type: Any, _exc_value: Any, traceback: Any) -> None: 

153 """Exits the asynchronous context manager. 

154 

155 This resets the `contextvar` to its previous state. 

156 

157 Args: 

158 _exc_type (Any): Exception type, if any (unused). 

159 _exc_value (Any): Exception value, if any (unused). 

160 traceback (Any): Traceback, if any (unused). 

161 

162 Returns: 

163 None: 

164 """ 

165 if self._token: 

166 _current_context.reset(self._token) 

167 self._token = None 

168 if os.getenv("VERBOSE_DI") and not os.getenv("BIJUXCLI_TEST_MODE"): 

169 self._log.log("debug", "Async context exited", extra={}) 

170 

171 @classmethod 

172 def current_data(cls) -> dict[str, Any]: 

173 """Returns the dictionary for the currently active CLI context. 

174 

175 This provides direct access to the data stored in the underlying 

176 `contextvar` for the current execution scope. 

177 

178 Returns: 

179 dict[str, Any]: The active context data dictionary. 

180 """ 

181 data = _current_context.get() 

182 if data is None: 

183 data = {} 

184 _current_context.set(data) 

185 return data 

186 

187 @classmethod 

188 def set_current_data(cls, data: dict[str, Any]) -> None: 

189 """Sets the dictionary for the currently active CLI context. 

190 

191 Args: 

192 data (dict[str, Any]): The data to use for the active context. 

193 

194 Returns: 

195 None: 

196 """ 

197 _current_context.set(data) 

198 

199 @classmethod 

200 @contextmanager 

201 def use_context(cls, data: dict[str, Any]) -> Iterator[None]: 

202 """Temporarily replaces the current context data within a `with` block. 

203 

204 Args: 

205 data (dict[str, Any]): The dictionary to use as the context for 

206 the duration of the `with` block. 

207 

208 Yields: 

209 None: 

210 """ 

211 token = _current_context.set(data) 

212 try: 

213 yield 

214 finally: 

215 _current_context.reset(token) 

216 

217 

218__all__ = ["Context"]