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

1import asyncio 

2import contextlib 

3import contextvars 

4import threading 

5from typing import Any, Dict, Union 

6 

7 

8class _CVar: 

9 """Storage utility for Local.""" 

10 

11 def __init__(self) -> None: 

12 self._data: "contextvars.ContextVar[Dict[str, Any]]" = contextvars.ContextVar( 

13 "asgiref.local" 

14 ) 

15 

16 def __getattr__(self, key): 

17 storage_object = self._data.get({}) 

18 try: 

19 return storage_object[key] 

20 except KeyError: 

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

22 

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

24 if key == "_data": 

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

26 

27 storage_object = self._data.get({}) 

28 storage_object[key] = value 

29 self._data.set(storage_object) 

30 

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

32 storage_object = self._data.get({}) 

33 if key in storage_object: 

34 del storage_object[key] 

35 self._data.set(storage_object) 

36 else: 

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

38 

39 

40class Local: 

41 """Local storage for async tasks. 

42 

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

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

45 

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

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

48 

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

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

51 `gather` and probably others). 

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

53 

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

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

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

57 and copies context across to that thread. 

58 

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

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

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

62 code (such as DB handles) to be kept stricly to their initial thread and 

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

64 

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

66 """ 

67 

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

69 self._thread_critical = thread_critical 

70 self._thread_lock = threading.RLock() 

71 

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

73 

74 if thread_critical: 

75 # Thread-local storage 

76 self._storage = threading.local() 

77 else: 

78 # Contextvar storage 

79 self._storage = _CVar() 

80 

81 @contextlib.contextmanager 

82 def _lock_storage(self): 

83 # Thread safe access to storage 

84 if self._thread_critical: 

85 try: 

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

87 # thread - will raise RuntimeError if there is 

88 # no current loop 

89 asyncio.get_running_loop() 

90 except RuntimeError: 

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

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

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

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

95 yield self._storage 

96 else: 

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

98 # local to this thread, but additionally should 

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

100 # the same async call stack) 

101 

102 # Ensure context exists in the current thread 

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

104 self._storage.cvar = _CVar() 

105 

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

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

108 # need any locks) 

109 yield self._storage.cvar 

110 else: 

111 # Lock for thread_critical=False as other threads 

112 # can access the exact same storage object 

113 with self._thread_lock: 

114 yield self._storage 

115 

116 def __getattr__(self, key): 

117 with self._lock_storage() as storage: 

118 return getattr(storage, key) 

119 

120 def __setattr__(self, key, value): 

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

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

123 with self._lock_storage() as storage: 

124 setattr(storage, key, value) 

125 

126 def __delattr__(self, key): 

127 with self._lock_storage() as storage: 

128 delattr(storage, key)