Merge pull request #98 from ikalnytskyi/chore/contextvars

Refactor picobox.contextvars

src/picobox/_scopes.py

@@ -4,6 +4,7 @@
 import contextvars as _contextvars
 import threading
 import typing as t
+import weakref
 
 
 class Scope(metaclass=abc.ABCMeta):
@@ -76,15 +77,22 @@ class contextvars(Scope):
     .. versionadded:: 2.1
     """
 
-    def __init__(self) -> None:
-        self._store: t.Dict[t.Hashable, _contextvars.ContextVar[t.Any]] = {}
+    _store_obj: (
+        "weakref.WeakKeyDictionary[Scope, t.Dict[t.Hashable, _contextvars.ContextVar[t.Any]]]"
+    )
+    _store_obj = weakref.WeakKeyDictionary()
 
-    def set(self, key: t.Hashable, value: t.Any) -> None:
+    @property
+    def _store(self) -> t.Dict[t.Hashable, _contextvars.ContextVar[t.Any]]:
         try:
-            var = self._store[key]
+            scope_store = self._store_obj[self]
         except KeyError:
-            var = self._store[key] = _contextvars.ContextVar("picobox")
-        var.set(value)
+            scope_store = self._store_obj[self] = {}
+        return scope_store
+
+    def set(self, key: t.Hashable, value: t.Any) -> None:
+        self._store[key] = _contextvars.ContextVar(str(key))
+        self._store[key].set(value)
 
     def get(self, key: t.Hashable) -> t.Any:
         try:

src/picobox/ext/asgiscopes.py

@@ -7,7 +7,7 @@
 import picobox
 
 if t.TYPE_CHECKING:
-    Store = weakref.WeakKeyDictionary[picobox.Scope, t.MutableMapping[t.Hashable, t.Any]]
+    Store = weakref.WeakKeyDictionary[picobox.Scope, t.Dict[t.Hashable, t.Any]]
     StoreCtxVar = contextvars.ContextVar[Store]
     ASGIScope = t.MutableMapping[str, t.Any]
     ASGIMessage = t.MutableMapping[str, t.Any]
@@ -65,7 +65,7 @@ class _asgiscope(picobox.Scope):
     _store_cvar: "StoreCtxVar"
 
     @property
-    def _store(self) -> t.MutableMapping[t.Hashable, t.Any]:
+    def _store(self) -> t.Dict[t.Hashable, t.Any]:
         try:
             store = self._store_cvar.get()
         except LookupError:

src/picobox/ext/wsgiscopes.py

@@ -9,7 +9,7 @@
 if t.TYPE_CHECKING:
     from _typeshed.wsgi import StartResponse, WSGIApplication, WSGIEnvironment
 
-    Store = weakref.WeakKeyDictionary[picobox.Scope, t.MutableMapping[t.Hashable, t.Any]]
+    Store = weakref.WeakKeyDictionary[picobox.Scope, t.Dict[t.Hashable, t.Any]]
     StoreCtxVar = contextvars.ContextVar[Store]
 
 
@@ -67,7 +67,7 @@ class _wsgiscope(picobox.Scope):
     _store_cvar: "StoreCtxVar"
 
     @property
-    def _store(self) -> t.MutableMapping[t.Hashable, t.Any]:
+    def _store(self) -> t.Dict[t.Hashable, t.Any]:
         try:
             store = self._store_cvar.get()
         except LookupError: