Dimension objects should stay singletons even when pickled

brian2/tests/test_units.py

@@ -202,6 +202,18 @@ def test_pickling():
         assert_equal(unpickled, q)
 
 
+@pytest.mark.codegen_independent
+def test_dimension_singletons():
+    # Make sure that Dimension objects are singletons, even when pickled
+    volt_dim = get_or_create_dimension((2, 1, -3, -1, 0, 0, 0))
+    assert volt.dim is volt_dim
+    import pickle
+    pickled_dim = pickle.dumps(volt_dim)
+    unpickled_dim = pickle.loads(pickled_dim)
+    assert unpickled_dim is volt_dim
+    assert unpickled_dim is volt.dim
+
+
 @pytest.mark.codegen_independent
 def test_str_repr():
     """

brian2/units/fundamentalunits.py

@@ -485,6 +485,10 @@ def __getstate__(self):
     def __setstate__(self, state):
         self._dims = state
 
+    def __reduce__(self):
+        # Make sure that unpickling Dimension objects does not bypass the singleton system
+        return (get_or_create_dimension, (self._dims, ))
+
     ### Dimension objects are singletons and deepcopy is therefore not necessary
     def __deepcopy__(self, memodict):
         return self