(partial) enable noop mode to work without an outer context too

amakelov · Jan 8, 2025 · b82e3a9 · b82e3a9
1 parent 201b809
commit b82e3a9
Show file tree

Hide file tree

Showing 2 changed files with 10 additions and 3 deletions.
diff --git a/mandala/storage.py b/mandala/storage.py
@@ -1207,13 +1207,15 @@ def __init__(self,):
         pass
 
     def __enter__(self) -> "Storage":
-        assert Context.current_context is not None, "A context must be active to use `noop`"
+        if Context.current_context is None:
+            return self
         storage = Context.current_context.storage
         res = storage(mode='noop')
         return res.__enter__()
 
     def __exit__(self, exc_type, exc_value, traceback):
-        assert Context.current_context is not None, "A context must be active to use `noop`"
+        if Context.current_context is None:
+            return
         storage = Context.current_context.storage
         res = storage(mode='run')
         return res.__exit__(exc_type, exc_value, traceback)

diff --git a/mandala/tests/test_modes.py b/mandala/tests/test_modes.py
@@ -67,4 +67,9 @@ def inc(x: int) -> int:
             z = inc(20)
         assert storage.mode == 'run'
         assert y == 22
-        assert z == 21
+        assert z == 21
+
+    # now, test it without a context
+    with noop():
+        x = inc(20)
+        assert x == 21