Context: initialize token at __init__

marshmallow-code · Jan 4, 2025 · e638af3 · e638af3
1 parent 318cae0
commit e638af3
Showing 1 changed file with 2 additions and 1 deletion.
diff --git a/src/marshmallow/experimental/context.py b/src/marshmallow/experimental/context.py
@@ -11,12 +11,13 @@
 class Context(contextlib.AbstractContextManager, typing.Generic[_T]):
     def __init__(self, context: _T) -> None:
         self.context = context
+        self.token: contextvars.Token | None = None
 
     def __enter__(self) -> None:
         self.token = _CURRENT_CONTEXT.set(self.context)
 
     def __exit__(self, *args, **kwargs) -> None:
-        _CURRENT_CONTEXT.reset(self.token)
+        _CURRENT_CONTEXT.reset(typing.cast(contextvars.Token, self.token))
 
     @classmethod
     def get(cls, default=...) -> _T: