Fix dataset pickling (#1218)

* Add regression test for dataset picklability * try just string? * try copyreg * changelog --------- Co-authored-by: Norman Rzepka <[email protected]>
scalableminds · Nov 25, 2024 · 5eb02eb · 5eb02eb
1 parent acc5eba
commit 5eb02eb
Show file tree

Hide file tree

Showing 3 changed files with 17 additions and 1 deletion.
diff --git a/webknossos/Changelog.md b/webknossos/Changelog.md
@@ -19,6 +19,7 @@ For upgrade instructions, please check the respective _Breaking Changes_ section
 ### Changed
 
 ### Fixed
+- Fixed pickling issue that has been introduced in 0.15.9. [#1218](https://github.com/scalableminds/webknossos-libs/pull/1218)
 
 
 ## [0.15.9](https://github.com/scalableminds/webknossos-libs/releases/tag/v0.15.9) - 2024-11-25
@@ -32,7 +33,7 @@ Removed the CZI installation extra from `pip install webknossos[all]` by default
 [Commits](https://github.com/scalableminds/webknossos-libs/compare/v0.15.7...v0.15.8)
 
 ### Changed
-- Fixes SSL certificate verification for remote datasets by adding CA certificates using `certifi`. [#1211](https://github.com/scalableminds/webknossos-libs/pull/1211)
+- Fixed SSL certificate verification for remote datasets by adding CA certificates using `certifi`. [#1211](https://github.com/scalableminds/webknossos-libs/pull/1211)
 
 
 ## [0.15.7](https://github.com/scalableminds/webknossos-libs/releases/tag/v0.15.7) - 2024-10-25

diff --git a/webknossos/tests/dataset/test_dataset_download_upload_remote.py b/webknossos/tests/dataset/test_dataset_download_upload_remote.py
@@ -1,3 +1,4 @@
+import pickle
 from pathlib import Path
 from tempfile import TemporaryDirectory
 from time import gmtime, strftime
@@ -73,6 +74,10 @@ def test_url_open_remote(
         data,
         sample_dataset.get_color_layers()[0].get_finest_mag().read(),
     )
+    assert set(pickle.loads(pickle.dumps(ds)).layers.keys()) == {
+        "color",
+        "segmentation",
+    }, "Dataset instances should be picklable."
 
 
 def test_remote_dataset(sample_dataset: wk.Dataset) -> None:

diff --git a/webknossos/webknossos/dataset/defaults.py b/webknossos/webknossos/dataset/defaults.py
@@ -1,3 +1,4 @@
+import copyreg
 import os
 import ssl
 
@@ -6,6 +7,15 @@
 from ..geometry import Vec3Int
 from .data_format import DataFormat
 
+
+def _save_sslcontext(
+    obj: ssl.SSLContext,
+) -> tuple[type[ssl.SSLContext], tuple[ssl._SSLMethod]]:
+    return obj.__class__, (obj.protocol,)
+
+
+copyreg.pickle(ssl.SSLContext, _save_sslcontext)
+
 DEFAULT_WKW_FILE_LEN = 32
 DEFAULT_CHUNK_SHAPE = Vec3Int.full(32)
 DEFAULT_DATA_FORMAT = (