Middleware API Demo

tdom/cvar_utils.py

@@ -0,0 +1,29 @@
+from contextvars import ContextVar, Token
+
+
+class ContextVarSetter:
+    """
+    Context manager for working with many context vars (instead of only 1).
+
+    This is meant to be created, used immediately and then discarded.
+
+    This allows for dynamically specifying a tuple of var / value pairs that
+    another part of the program can use to wrap some called code without knowing
+    anything about either.
+    """
+
+    context_values: tuple[tuple[ContextVar, object], ...]  # Cvar / value pair.
+    tokens: tuple[Token, ...]
+
+    def __init__(self, context_values=()):
+        self.context_values = context_values
+        self.tokens = ()
+
+    def __enter__(self):
+        """Set every given context var to its paired value."""
+        self.tokens = tuple(var.set(val) for var, val in self.context_values)
+
+    def __exit__(self, exc_type, exc_value, traceback):
+        """Reset every given context var."""
+        for idx, var_value in enumerate(self.context_values):
+            var_value[0].reset(self.tokens[idx])

tdom/cvar_utils_test.py

@@ -0,0 +1,72 @@
+import string
+from contextvars import ContextVar
+
+from .cvar_utils import ContextVarSetter
+
+CtxStr = ContextVar[str]("CtxStr", default="default")
+CtxInt = ContextVar[int]("CtxInt", default=0)
+
+
+def _assert_ctx(ctx_str: str = "default", ctx_int: int = 0):
+    assert CtxStr.get() == ctx_str
+    assert CtxInt.get() == ctx_int
+
+
+def test_set():
+    _assert_ctx()
+    with ContextVarSetter(
+        context_values=(
+            (CtxStr, "new"),
+            (CtxInt, 1),
+        )
+    ):
+        _assert_ctx("new", 1)
+    _assert_ctx()
+
+
+def test_nest():
+    _assert_ctx()
+    with ContextVarSetter(
+        context_values=(
+            (CtxStr, "new"),
+            (CtxInt, 1),
+        )
+    ):
+        _assert_ctx("new", 1)
+        with ContextVarSetter(
+            context_values=(
+                (CtxStr, "again"),
+                (CtxInt, 2),
+            )
+        ):
+            _assert_ctx("again", 2)
+        _assert_ctx("new", 1)
+    _assert_ctx()
+
+
+def test_reps():
+    _assert_ctx()
+    for index, leter in enumerate(string.ascii_lowercase):
+        with ContextVarSetter(
+            context_values=(
+                (CtxStr, leter),
+                (CtxInt, index),
+            )
+        ):
+            _assert_ctx(leter, index)
+        _assert_ctx()
+
+
+def test_empty():
+    _assert_ctx()
+    with ContextVarSetter(context_values=()):
+        # DO NOTHING BUT NOT AN ERROR
+        _assert_ctx()
+    _assert_ctx()
+
+
+def test_one():
+    _assert_ctx()
+    with ContextVarSetter(context_values=((CtxStr, "new"),)):
+        _assert_ctx("new")
+    _assert_ctx()