Make Views picklable (#469)

* make Views picklable * update changelog * Update webknossos/tests/test_dataset.py Co-authored-by: Jonathan Striebel <[email protected]> * Update webknossos/tests/test_dataset.py Co-authored-by: Jonathan Striebel <[email protected]> Co-authored-by: Jonathan Striebel <[email protected]>
scalableminds · Oct 28, 2021 · 9a9e8c0 · 9a9e8c0
1 parent cbd1bed
commit 9a9e8c0
Show file tree

Hide file tree

Showing 3 changed files with 33 additions and 1 deletion.
diff --git a/webknossos/Changelog.md b/webknossos/Changelog.md
@@ -20,6 +20,7 @@ For upgrade instructions, please check the respective *Breaking Changes* section
 ### Changed
 - `View`s now always open the `wkw.Dataset` lazily. All explicit calls to `View.open()` and `View.close()` must be removed. [#448](https://github.com/scalableminds/webknossos-libs/pull/448)
 ### Fixed
+- Make Views picklable. We now ignore the file handle when we pickle Views. [#469](https://github.com/scalableminds/webknossos-libs/pull/469)
 
 
 ## [0.8.18](https://github.com/scalableminds/webknossos-cuber/releases/tag/v0.8.18) - 2021-10-18

diff --git a/webknossos/tests/test_dataset.py b/webknossos/tests/test_dataset.py
@@ -1,6 +1,7 @@
 import itertools
 import json
 import os
+import pickle
 import warnings
 from os.path import join
 from pathlib import Path
@@ -1957,3 +1958,24 @@ def test_add_layer_like(tmp_path: Path) -> None:
     )
 
     assure_exported_properties(ds)
+
+
+def test_pickle_view(tmp_path: Path) -> None:
+    ds = Dataset.create(tmp_path / "ds", scale=(1, 1, 1))
+    mag1 = ds.add_layer("color", LayerCategories.COLOR_TYPE).add_mag(1)
+
+    assert mag1._cached_wkw_dataset is None
+    data_to_write = (np.random.rand(1, 10, 10, 10) * 255).astype(np.uint8)
+    mag1.write(data_to_write)
+    assert mag1._cached_wkw_dataset is not None
+
+    pickle.dump(mag1, open(str(tmp_path / "save.p"), "wb"))
+    pickled_mag1 = pickle.load(open(str(tmp_path / "save.p"), "rb"))
+
+    # Make sure that the pickled mag can still read data
+    assert pickled_mag1._cached_wkw_dataset is None
+    assert np.array_equal(data_to_write, pickled_mag1.read())
+    assert pickled_mag1._cached_wkw_dataset is not None
+
+    # Make sure that the attributes of the MagView (not View) still exist
+    assert pickled_mag1.layer is not None
diff --git a/webknossos/webknossos/dataset/view.py b/webknossos/webknossos/dataset/view.py
@@ -2,7 +2,7 @@
 import warnings
 from pathlib import Path
 from types import TracebackType
-from typing import TYPE_CHECKING, Callable, Optional, Tuple, Type, Union
+from typing import TYPE_CHECKING, Any, Callable, Dict, Optional, Tuple, Type, Union
 
 import cluster_tools
 import numpy as np
@@ -585,6 +585,15 @@ def _wkw_dataset(self) -> None:
     def __del__(self) -> None:
         del self._cached_wkw_dataset
 
+    def __getstate__(self) -> Dict[str, Any]:
+        d = dict(self.__dict__)
+        del d["_cached_wkw_dataset"]
+        return d
+
+    def __setstate__(self, d: Dict[str, Any]) -> None:
+        d["_cached_wkw_dataset"] = None
+        self.__dict__ = d
+
 
 def _assert_positive_dimensions(offset: Vec3Int, size: Vec3Int) -> None:
     if any(x < 0 for x in offset):