xgc: Support nested gc_disabled() contexts

davidfstr · Sep 8, 2023 · 1afd5fa · 1afd5fa
1 parent 48731b4
commit 1afd5fa
Showing 1 changed file with 5 additions and 3 deletions.
diff --git a/src/crystal/util/xgc.py b/src/crystal/util/xgc.py
@@ -45,10 +45,12 @@ def gc_disabled() -> Iterator[None]:
     Useful when allocating a very large number of long-lived objects,
     to avoid many useless runs of the garbage collector.
     """
-    assert gc.isenabled()
-    gc.disable()
+    was_enabled = gc.isenabled()
+    if was_enabled:
+        gc.disable()
     try:
         yield
     finally:
         assert not gc.isenabled()
-        gc.enable()
+        if was_enabled:
+            gc.enable()