Merge 345a676 into bbfeabe

qctoolkit/serialization.py

@@ -18,6 +18,7 @@
 import weakref
 import warnings
 import gc
+from contextlib import contextmanager
 
 from qctoolkit.utils.types import DocStringABCMeta
 
@@ -320,6 +321,10 @@ def __new__(mcs, name, bases, dct):
 default_pulse_registry = weakref.WeakValueDictionary()
 
 
+def get_default_pulse_registry() -> Union[weakref.WeakKeyDictionary, 'PulseStorage']:
+    return default_pulse_registry
+
+
 class Serializable(metaclass=SerializableMeta):
     """Any object that can be converted into a serialized representation for storage and back.
 
@@ -697,6 +702,20 @@ def overwrite(self, identifier: str, serializable: Serializable) -> None:
     def clear(self) -> None:
         self._temporary_storage.clear()
 
+    @contextmanager
+    def as_default_registry(self) -> Any:
+        global default_pulse_registry
+        previous_registry = default_pulse_registry
+        default_pulse_registry = self
+        try:
+            yield self
+        finally:
+            default_pulse_registry = previous_registry
+
+    def set_to_default_registry(self) -> None:
+        global default_pulse_registry
+        default_pulse_registry = self
+
 
 class JSONSerializableDecoder(json.JSONDecoder):
 

tests/serialization_tests.py

@@ -13,7 +13,8 @@
 
 from qctoolkit.serialization import FilesystemBackend, CachingBackend, Serializable, JSONSerializableEncoder,\
     ZipFileBackend, AnonymousSerializable, DictBackend, PulseStorage, JSONSerializableDecoder, Serializer,\
-    default_pulse_registry
+    get_default_pulse_registry
+
 from qctoolkit.expressions import ExpressionScalar
 
 from tests.serialization_dummies import DummyStorageBackend
@@ -119,7 +120,7 @@ def test_serialization_and_deserialization(self):
         self.assert_equal_instance(instance, other_instance)
 
         self.assertIs(registry['blub'], instance)
-        self.assertIs(default_pulse_registry['blub'], other_instance)
+        self.assertIs(get_default_pulse_registry()['blub'], other_instance)
 
     def test_duplication_error(self):
         registry = dict()
@@ -554,6 +555,23 @@ def test_clear(self):
 
         self.assertFalse(self.storage.temporary_storage)
 
+    def test_as_default_registry(self) -> None:
+        prev_reg = get_default_pulse_registry()
+        pulse_storage = PulseStorage(DummyStorageBackend())
+        with pulse_storage.as_default_registry():
+            self.assertIs(get_default_pulse_registry(), pulse_storage)
+        self.assertIs(get_default_pulse_registry(), prev_reg)
+
+    def test_set_to_default_registry(self) -> None:
+        pulse_storage = PulseStorage(DummyStorageBackend())
+        previous_default_registry = get_default_pulse_registry()
+        try:
+            pulse_storage.set_to_default_registry()
+            self.assertIs(get_default_pulse_registry(), pulse_storage)
+        finally:
+            import qctoolkit.serialization
+            qctoolkit.serialization.default_pulse_registry = previous_default_registry
+
     def test_beautified_json(self) -> None:
         data = {'e': 89, 'b': 151, 'c': 123515, 'a': 123, 'h': 2415}
         template = DummySerializable(data=data)