Simplify code leveraging the fact that ContextVar is automatically thread-local

Author: spanezz

asgiref/local.py

@@ -38,7 +38,7 @@ def __delattr__(self, key: str) -> None:
             raise AttributeError(f"{self!r} object has no attribute {key!r}")
 
 
-class Local:
+def Local(thread_critical: bool = False) -> threading.local | _CVar:
     """Local storage for async tasks.
 
     This is a namespace object (similar to `threading.local`) where data is
@@ -65,65 +65,7 @@ class Local:
 
     Unlike plain `contextvars` objects, this utility is threadsafe.
     """
-
-    def __init__(self, thread_critical: bool = False) -> None:
-        self._thread_critical = thread_critical
-        self._thread_lock = threading.RLock()
-
-        self._storage: "Union[threading.local, _CVar]"
-
-        if thread_critical:
-            # Thread-local storage
-            self._storage = threading.local()
-        else:
-            # Contextvar storage
-            self._storage = _CVar()
-
-    @contextlib.contextmanager
-    def _lock_storage(self):
-        # Thread safe access to storage
-        if self._thread_critical:
-            try:
-                # this is a test for are we in a async or sync
-                # thread - will raise RuntimeError if there is
-                # no current loop
-                asyncio.get_running_loop()
-            except RuntimeError:
-                # We are in a sync thread, the storage is
-                # just the plain thread local (i.e, "global within
-                # this thread" - it doesn't matter where you are
-                # in a call stack you see the same storage)
-                yield self._storage
-            else:
-                # We are in an async thread - storage is still
-                # local to this thread, but additionally should
-                # behave like a context var (is only visible with
-                # the same async call stack)
-
-                # Ensure context exists in the current thread
-                if not hasattr(self._storage, "cvar"):
-                    self._storage.cvar = _CVar()
-
-                # self._storage is a thread local, so the members
-                # can't be accessed in another thread (we don't
-                # need any locks)
-                yield self._storage.cvar
-        else:
-            # Lock for thread_critical=False as other threads
-            # can access the exact same storage object
-            with self._thread_lock:
-                yield self._storage
-
-    def __getattr__(self, key):
-        with self._lock_storage() as storage:
-            return getattr(storage, key)
-
-    def __setattr__(self, key, value):
-        if key in ("_local", "_storage", "_thread_critical", "_thread_lock"):
-            return super().__setattr__(key, value)
-        with self._lock_storage() as storage:
-            setattr(storage, key, value)
-
-    def __delattr__(self, key):
-        with self._lock_storage() as storage:
-            delattr(storage, key)
+    if thread_critical:
+        return threading.local()
+    else:
+        return _CVar()