fix pickle

xsuite · Aug 16, 2024 · 341319a · 341319a
1 parent ebb22a4
commit 341319a
Show file tree

Hide file tree

Showing 2 changed files with 12 additions and 0 deletions.
diff --git a/xobjects/context.py b/xobjects/context.py
@@ -222,6 +222,15 @@ def new_buffer(self, capacity=1048576):
         self._allocations += 1
         return buf
 
+    def __getstate__(self):
+        state = self.__dict__
+        del state["_buffers"]
+        return state
+
+    def __setstate__(self, state):
+        self.__dict__.update(state)
+        self._buffers = weakref.WeakSet()
+
     @property
     def buffers(self):
         return list(self._buffers)

diff --git a/xobjects/context_cpu.py b/xobjects/context_cpu.py
@@ -10,6 +10,7 @@
 import uuid
 from pathlib import Path
 from typing import Callable, Dict, List, Sequence, Tuple
+import weakref
 
 from .general import _print
 
@@ -645,10 +646,12 @@ def openmp_enabled(self):
     def __getstate__(self):
         state = self.__dict__.copy()
         state["_kernels"] = {}
+        del state["_buffers"]
         return state
 
     def __setstate__(self, state):
         self.__dict__.update(state)
+        self._buffers = weakref.WeakSet()
 
 
 class BufferByteArray(XBuffer):