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

1# SPDX-License-Identifier: Apache-2.0 

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 `ExecutionContext`. 

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 

18from typing import Any, TypeVar 

19 

20from injector import inject 

21 

22from bijux_cli.core.contracts import ExecutionContext 

23from bijux_cli.core.di import DIContainer 

24 

25T = TypeVar("T") 

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

27 "current_context", default=None 

28) 

29 

30 

31class Context(ExecutionContext): 

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

33 

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

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

36 both a synchronous and asynchronous context manager. 

37 

38 Attributes: 

39 _di (DIContainer): The dependency injection container. 

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

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

42 """ 

43 

44 @inject 

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

46 """Initializes a new Context instance. 

47 

48 Args: 

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

50 resolve the logging service. 

51 """ 

52 self._di = di 

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

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

55 

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

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

58 

59 Args: 

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

61 value (Any): The value to store. 

62 

63 Returns: 

64 None: 

65 """ 

66 self._data[key] = value 

67 

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

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

70 

71 Args: 

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

73 

74 Returns: 

75 Any: The value associated with the key. 

76 

77 Raises: 

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

79 """ 

80 if key not in self._data: 

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

82 return self._data[key] 

83 

84 def clear(self) -> None: 

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

86 self._data.clear() 

87 

88 def __enter__(self) -> Context: 

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

90 

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

92 

93 Returns: 

94 Context: The context instance itself. 

95 """ 

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

97 return self 

98 

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

100 """Exits the synchronous context manager. 

101 

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

103 

104 Args: 

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

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

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

108 

109 Returns: 

110 None: 

111 """ 

112 if self._token: 112 ↛ exitline 112 didn't return from function '__exit__' because the condition on line 112 was always true

113 _current_context.reset(self._token) 

114 self._token = None 

115 

116 async def __aenter__(self) -> Context: 

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

118 

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

120 

121 Returns: 

122 Context: The context instance itself. 

123 """ 

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

125 return self 

126 

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

128 """Exits the asynchronous context manager. 

129 

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

131 

132 Args: 

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

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

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

136 

137 Returns: 

138 None: 

139 """ 

140 if self._token: 140 ↛ exitline 140 didn't return from function '__aexit__' because the condition on line 140 was always true

141 _current_context.reset(self._token) 

142 self._token = None 

143 

144 @classmethod 

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

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

147 

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

149 `contextvar` for the current execution scope. 

150 

151 Returns: 

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

153 """ 

154 data = _current_context.get() 

155 if data is None: 

156 data = {} 

157 _current_context.set(data) 

158 return data 

159 

160 @classmethod 

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

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

163 

164 Args: 

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

166 

167 Returns: 

168 None: 

169 """ 

170 _current_context.set(data) 

171 

172 @classmethod 

173 @contextmanager 

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

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

176 

177 Args: 

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

179 the duration of the `with` block. 

180 

181 Yields: 

182 None: 

183 """ 

184 token = _current_context.set(data) 

185 try: 

186 yield 

187 finally: 

188 _current_context.reset(token) 

189 

190 

191__all__ = ["Context"]