Coverage for /pythoncovmergedfiles/medio/medio/usr/local/lib/python3.11/site-packages/asgiref/local.py: 65%

1import asyncio 

2import contextlib 

3import contextvars 

4import threading 

5from typing import Any, Union 

6 

7 

8class _Storage: 

9 """Thread-tagged storage for a non-thread-critical ``Local``. 

10 

11 The data is tagged with the identity of the thread that owns it. This lets 

12 ``_CVar`` ignore data that leaked into an unrelated thread. 

13 

14 Python 3.14 added ``sys.flags.thread_inherit_context``, which is enabled by 

15 default on free-threaded builds. When set, a new thread starts with a copy 

16 of the spawning thread's context instead of an empty one, so the contextvar 

17 backing a ``Local`` would otherwise be visible in any thread spawned from 

18 one that had set it -- breaking the documented "thread-local in sync 

19 threads" behaviour. asgiref re-homes the storage to the current thread at 

20 the points where it *intentionally* moves work between threads (see 

21 ``asgiref.sync._restore_context``); data merely inherited by an unrelated 

22 thread is never re-homed and so stays isolated. 

23 """ 

24 

25 __slots__ = ("thread_id", "data") 

26 

27 def __init__(self, thread_id: int, data: dict[str, Any]) -> None: 

28 self.thread_id = thread_id 

29 self.data = data 

30 

31 

32def _rehome(storage: "_Storage") -> "_Storage": 

33 """Return a copy of *storage* owned by the current thread.""" 

34 return _Storage(threading.get_ident(), storage.data) 

35 

36 

37class _CVar: 

38 """Storage utility for Local.""" 

39 

40 def __init__(self) -> None: 

41 self._data: "contextvars.ContextVar[_Storage]" = contextvars.ContextVar( 

42 "asgiref.local" 

43 ) 

44 

45 def _storage(self) -> "_Storage": 

46 # Only return storage that belongs to the current thread. Storage with 

47 # a different thread id was inherited by this thread (rather than 

48 # intentionally moved here by asgiref) and must not be visible. 

49 storage = self._data.get(None) 

50 if storage is None or storage.thread_id != threading.get_ident(): 

51 return _Storage(threading.get_ident(), {}) 

52 return storage 

53 

54 def __getattr__(self, key): 

55 try: 

56 return self._storage().data[key] 

57 except KeyError: 

58 raise AttributeError(f"{self!r} object has no attribute {key!r}") 

59 

60 def __setattr__(self, key: str, value: Any) -> None: 

61 if key == "_data": 

62 return super().__setattr__(key, value) 

63 

64 data = self._storage().data.copy() 

65 data[key] = value 

66 self._data.set(_Storage(threading.get_ident(), data)) 

67 

68 def __delattr__(self, key: str) -> None: 

69 data = self._storage().data.copy() 

70 if key in data: 

71 del data[key] 

72 self._data.set(_Storage(threading.get_ident(), data)) 

73 else: 

74 raise AttributeError(f"{self!r} object has no attribute {key!r}") 

75 

76 

77class Local: 

78 """Local storage for async tasks. 

79 

80 This is a namespace object (similar to `threading.local`) where data is 

81 also local to the current async task (if there is one). 

82 

83 In async threads, local means in the same sense as the `contextvars` 

84 module - i.e. a value set in an async frame will be visible: 

85 

86 - to other async code `await`-ed from this frame. 

87 - to tasks spawned using `asyncio` utilities (`create_task`, `wait_for`, 

88 `gather` and probably others). 

89 - to code scheduled in a sync thread using `sync_to_async` 

90 

91 In "sync" threads (a thread with no async event loop running), the 

92 data is thread-local, but additionally shared with async code executed 

93 via the `async_to_sync` utility, which schedules async code in a new thread 

94 and copies context across to that thread. 

95 

96 If `thread_critical` is True, then the local will only be visible per-thread, 

97 behaving exactly like `threading.local` if the thread is sync, and as 

98 `contextvars` if the thread is async. This allows genuinely thread-sensitive 

99 code (such as DB handles) to be kept strictly to their initial thread and 

100 disable the sharing across `sync_to_async` and `async_to_sync` wrapped calls. 

101 

102 Unlike plain `contextvars` objects, this utility is threadsafe. 

103 """ 

104 

105 def __init__(self, thread_critical: bool = False) -> None: 

106 self._thread_critical = thread_critical 

107 self._thread_lock = threading.RLock() 

108 

109 self._storage: "Union[threading.local, _CVar]" 

110 

111 if thread_critical: 

112 # Thread-local storage 

113 self._storage = threading.local() 

114 else: 

115 # Contextvar storage 

116 self._storage = _CVar() 

117 

118 @contextlib.contextmanager 

119 def _lock_storage(self): 

120 # Thread safe access to storage 

121 if self._thread_critical: 

122 is_async = True 

123 try: 

124 # this is a test for are we in a async or sync 

125 # thread - will raise RuntimeError if there is 

126 # no current loop 

127 asyncio.get_running_loop() 

128 except RuntimeError: 

129 is_async = False 

130 if not is_async: 

131 # We are in a sync thread, the storage is 

132 # just the plain thread local (i.e, "global within 

133 # this thread" - it doesn't matter where you are 

134 # in a call stack you see the same storage) 

135 yield self._storage 

136 else: 

137 # We are in an async thread - storage is still 

138 # local to this thread, but additionally should 

139 # behave like a context var (is only visible with 

140 # the same async call stack) 

141 

142 # Ensure context exists in the current thread 

143 if not hasattr(self._storage, "cvar"): 

144 self._storage.cvar = _CVar() 

145 

146 # self._storage is a thread local, so the members 

147 # can't be accessed in another thread (we don't 

148 # need any locks) 

149 yield self._storage.cvar 

150 else: 

151 # Lock for thread_critical=False as other threads 

152 # can access the exact same storage object 

153 with self._thread_lock: 

154 yield self._storage 

155 

156 def __getattr__(self, key): 

157 with self._lock_storage() as storage: 

158 return getattr(storage, key) 

159 

160 def __setattr__(self, key, value): 

161 if key in ("_local", "_storage", "_thread_critical", "_thread_lock"): 

162 return super().__setattr__(key, value) 

163 with self._lock_storage() as storage: 

164 setattr(storage, key, value) 

165 

166 def __delattr__(self, key): 

167 with self._lock_storage() as storage: 

168 delattr(storage, key)