diff --git a/xobjects/context.py b/xobjects/context.py index 776e34f..506fbce 100644 --- 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 index d716ed8..a03b443 100644 --- 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):