Support for serializing CArray too

Author: FrancescAlted

src/blosc2/_msgpack_utils.py

@@ -16,6 +16,48 @@
 # reconstructed with ``blosc2.from_cframe()``.  Keep this stable for backward
 # compatibility with persisted msgpack payloads produced by this package.
 _BLOSC2_EXT_CODE = 42
+# Reserve code 43 for structured Blosc2 reference objects that are not naturally
+# serialized as CFrames.  The payload is a msgpack-encoded mapping with a
+# stable ``kind`` and ``version`` envelope.
+_BLOSC2_STRUCTURED_EXT_CODE = 43
+_BLOSC2_STRUCTURED_VERSION = 1
+
+
+def _encode_structured_reference(obj):
+    import blosc2
+
+    if isinstance(obj, blosc2.C2Array):
+        payload = {
+            "kind": "c2array",
+            "version": _BLOSC2_STRUCTURED_VERSION,
+            "path": obj.path,
+            "urlbase": obj.urlbase,
+        }
+        return ExtType(_BLOSC2_STRUCTURED_EXT_CODE, packb(payload, use_bin_type=True))
+    return None
+
+
+def _decode_structured_reference(data):
+    import blosc2
+
+    payload = unpackb(data)
+    if not isinstance(payload, dict):
+        raise TypeError("Structured Blosc2 msgpack payload must decode to a mapping")
+
+    version = payload.get("version")
+    if version != _BLOSC2_STRUCTURED_VERSION:
+        raise ValueError(f"Unsupported structured Blosc2 msgpack payload version: {version!r}")
+
+    kind = payload.get("kind")
+    if kind == "c2array":
+        path = payload.get("path")
+        if not isinstance(path, str):
+            raise TypeError("Structured C2Array msgpack payload requires a string 'path'")
+        urlbase = payload.get("urlbase")
+        if urlbase is not None and not isinstance(urlbase, str):
+            raise TypeError("Structured C2Array msgpack payload requires 'urlbase' to be a string or None")
+        return blosc2.C2Array(path, urlbase=urlbase)
+    raise ValueError(f"Unsupported structured Blosc2 msgpack payload kind: {kind!r}")
 
 
 def _encode_msgpack_ext(obj):
@@ -32,6 +74,9 @@ def _encode_msgpack_ext(obj):
         ),
     ):
         return ExtType(_BLOSC2_EXT_CODE, obj.to_cframe())
+    structured = _encode_structured_reference(obj)
+    if structured is not None:
+        return structured
     return blosc2_ext.encode_tuple(obj)
 
 
@@ -50,6 +95,8 @@ def _decode_msgpack_ext(code, data):
 
     if code == _BLOSC2_EXT_CODE:
         return blosc2.from_cframe(data, copy=True)
+    if code == _BLOSC2_STRUCTURED_EXT_CODE:
+        return _decode_structured_reference(data)
     return ExtType(code, data)
 
 

tests/test_batch_store.py

@@ -9,6 +9,7 @@
 import pytest
 
 import blosc2
+import blosc2.c2array as blosc2_c2array
 from blosc2._msgpack_utils import msgpack_packb, msgpack_unpackb
 
 BATCHES = [
@@ -46,6 +47,14 @@ def _make_nested_blosc2_objects():
     return ndarray, schunk, nested_vlarray, nested_batchstore, estore
 
 
+def _make_c2array(monkeypatch, path="@public/examples/ds-1d.b2nd", urlbase="https://cat2.cloud/demo/"):
+    def fake_info(path_, urlbase_, params=None, headers=None, model=None, auth_token=None):
+        return {"schunk": {"cparams": dict(blosc2.cparams_dflts)}}
+
+    monkeypatch.setattr(blosc2_c2array, "info", fake_info)
+    return blosc2.C2Array(path, urlbase=urlbase)
+
+
 @pytest.mark.parametrize(
     ("contiguous", "urlpath"),
     [
@@ -225,6 +234,62 @@ def test_batchstore_msgpack_supports_blosc2_objects():
     assert np.array_equal(restored[4]["/node"][:], estore["/node"][:])
 
 
+def test_msgpack_supports_c2array(monkeypatch):
+    c2array = _make_c2array(monkeypatch)
+
+    payload = msgpack_packb({"remote": c2array})
+    restored = msgpack_unpackb(payload)
+
+    assert isinstance(restored["remote"], blosc2.C2Array)
+    assert restored["remote"].path == c2array.path
+    assert restored["remote"].urlbase == c2array.urlbase
+    assert restored["remote"].auth_token is None
+
+
+def test_batchstore_msgpack_supports_c2array(monkeypatch):
+    c2array = _make_c2array(monkeypatch)
+
+    barray = blosc2.BatchStore(items_per_block=2)
+    barray.append([c2array])
+
+    restored = barray[0][0]
+
+    assert isinstance(restored, blosc2.C2Array)
+    assert restored.path == c2array.path
+    assert restored.urlbase == c2array.urlbase
+    assert restored.auth_token is None
+
+
+@pytest.mark.network
+def test_msgpack_roundtrip_c2array_network(cat2_context):
+    path = "@public/expr/ds-1-2-linspace-float64-b2-(5,)d.b2nd"
+    original = blosc2.C2Array(path)
+
+    payload = msgpack_packb({"remote": original})
+    restored = msgpack_unpackb(payload)["remote"]
+
+    assert isinstance(restored, blosc2.C2Array)
+    assert restored.path == original.path
+    assert restored.urlbase == original.urlbase
+    np.testing.assert_allclose(restored[:], original[:])
+
+
+@pytest.mark.network
+def test_batchstore_msgpack_roundtrip_c2array_network(cat2_context):
+    path = "@public/expr/ds-1-2-linspace-float64-b2-(5,)d.b2nd"
+    original = blosc2.C2Array(path)
+
+    barray = blosc2.BatchStore(items_per_block=2)
+    barray.append([original])
+
+    restored = barray[0][0]
+
+    assert isinstance(restored, blosc2.C2Array)
+    assert restored.path == original.path
+    assert restored.urlbase == original.urlbase
+    np.testing.assert_allclose(restored[:], original[:])
+
+
 def test_batchstore_info():
     barray = blosc2.BatchStore()
     barray.extend(BATCHES)

tests/test_vlarray.py

@@ -9,6 +9,7 @@
 import pytest
 
 import blosc2
+import blosc2.c2array as blosc2_c2array
 
 VALUES = [
     b"bytes\x00payload",
@@ -45,6 +46,14 @@ def _make_nested_blosc2_objects():
     return ndarray, schunk, nested_vlarray, nested_batchstore, estore
 
 
+def _make_c2array(monkeypatch, path="@public/examples/ds-1d.b2nd", urlbase="https://cat2.cloud/demo/"):
+    def fake_info(path_, urlbase_, params=None, headers=None, model=None, auth_token=None):
+        return {"schunk": {"cparams": dict(blosc2.cparams_dflts)}}
+
+    monkeypatch.setattr(blosc2_c2array, "info", fake_info)
+    return blosc2.C2Array(path, urlbase=urlbase)
+
+
 @pytest.mark.parametrize(
     ("contiguous", "urlpath"),
     [
@@ -169,6 +178,36 @@ def test_vlarray_msgpack_supports_blosc2_objects():
     assert np.array_equal(restored["estore"]["/node"][:], estore["/node"][:])
 
 
+def test_vlarray_msgpack_supports_c2array(monkeypatch):
+    c2array = _make_c2array(monkeypatch)
+
+    vlarray = blosc2.VLArray()
+    vlarray.append(c2array)
+
+    restored = vlarray[0]
+
+    assert isinstance(restored, blosc2.C2Array)
+    assert restored.path == c2array.path
+    assert restored.urlbase == c2array.urlbase
+    assert restored.auth_token is None
+
+
+@pytest.mark.network
+def test_vlarray_msgpack_roundtrip_c2array_network(cat2_context):
+    path = "@public/expr/ds-1-2-linspace-float64-b2-(5,)d.b2nd"
+    original = blosc2.C2Array(path)
+
+    vlarray = blosc2.VLArray()
+    vlarray.append(original)
+
+    restored = vlarray[0]
+
+    assert isinstance(restored, blosc2.C2Array)
+    assert restored.path == original.path
+    assert restored.urlbase == original.urlbase
+    np.testing.assert_allclose(restored[:], original[:])
+
+
 def test_vlarray_info():
     vlarray = blosc2.VLArray()
     vlarray.extend(VALUES)