Add json+binary codec

Author: benjeffery

python/CHANGELOG.rst

@@ -107,6 +107,9 @@ Maintenance release.
   also around 10% faster.
   (:user:`benjeffery`, :pr:`3313`, :pr:`3317`, :issue:`1896`)
 
+- Add ``json+binary`` metadata codec that allows storing binary data
+  alongside JSON metadata. (:user:`benjeffery`, :pr:`3306`)
+
 **Bugfixes**
 
 - In some tables with mutations out-of-order ``TableCollection.sort`` did not re-order

python/tests/test_metadata.py

@@ -627,6 +627,133 @@ def test_zero_length(self):
         assert ms.decode_row(b"") == {}
 
 
+class TestJSONBinaryCodec:
+    def test_encode_requires_binary(self):
+        ms = tskit.MetadataSchema({"codec": "json+binary"})
+        with pytest.raises(
+            exceptions.MetadataEncodingError,
+            match="requires top-level '_binary' bytes-like value",
+        ):
+            ms.validate_and_encode_row({})
+
+    def test_zero_length_blob(self):
+        ms = tskit.MetadataSchema({"codec": "json+binary"})
+        encoded = ms.validate_and_encode_row({"_binary": b""})
+        decoded = ms.decode_row(encoded)
+        assert isinstance(decoded["_binary"], memoryview)
+        assert len(decoded["_binary"]) == 0
+        # JSON portion was empty
+        assert set(decoded.keys()) == {"_binary"}
+
+    def test_round_trip_with_blob_and_json(self):
+        ms = tskit.MetadataSchema({"codec": "json+binary"})
+        blob = b"\x00\x01\x02hello"
+        row = {"label": "alpha", "count": 7, "_binary": blob}
+        encoded = ms.validate_and_encode_row(row)
+        out = ms.decode_row(encoded)
+        assert out["label"] == "alpha"
+        assert out["count"] == 7
+        assert isinstance(out["_binary"], memoryview)
+        assert out["_binary"].tobytes() == blob
+
+    def test_decode_without_magic_errors(self):
+        ms = tskit.MetadataSchema({"codec": "json+binary"})
+        # Plain JSON is not acceptable for this codec
+        with pytest.raises(ValueError, match="missing magic header"):
+            ms.decode_row(b"{}")
+
+    def test_simple_default(self):
+        schema = {
+            "codec": "json+binary",
+            "type": "object",
+            "properties": {"number": {"type": "number", "default": 5}},
+        }
+        ms = tskit.MetadataSchema(schema)
+        # With json+binary, we need to provide _binary even for empty metadata
+        assert ms.decode_row(ms.validate_and_encode_row({"_binary": b""})) == {
+            "number": 5,
+            "_binary": memoryview(b""),
+        }
+        assert ms.decode_row(
+            ms.validate_and_encode_row({"_binary": b"", "number": 42})
+        ) == {"number": 42, "_binary": memoryview(b"")}
+
+    def test_nested_default_error(self):
+        schema = {
+            "codec": "json+binary",
+            "type": "object",
+            "properties": {
+                "obj": {
+                    "type": "object",
+                    "properties": {
+                        "nested_obj_no_default": {
+                            "type": "object",
+                            "properties": {},
+                        },
+                        "nested_obj": {
+                            "type": "object",
+                            "properties": {},
+                            "default": {"foo": "bar"},
+                        },
+                    },
+                }
+            },
+        }
+        with pytest.raises(
+            tskit.MetadataSchemaValidationError,
+            match="Defaults can only be specified at the top level for JSON codec",
+        ):
+            tskit.MetadataSchema(schema)
+
+    def test_bad_type_error(self):
+        ms = tskit.MetadataSchema({"codec": "json+binary"})
+        # json+binary first checks for _binary key, so we need a dict with _binary
+        # but other fields that can't be JSON encoded
+        with pytest.raises(
+            exceptions.MetadataEncodingError,
+            match="Could not encode metadata of type TableCollection",
+        ):
+            ms.validate_and_encode_row(
+                {"_binary": b"", "bad_field": tskit.TableCollection(1)}
+            )
+
+    def test_skip_validation(self):
+        ms = tskit.MetadataSchema({"codec": "json+binary"})
+        assert ms._bypass_validation
+        with patch.object(ms, "_validate_row", return_value=True) as mocked_validate:
+            ms.validate_and_encode_row({"_binary": b""})
+            assert mocked_validate.call_count == 0
+
+    def test_dont_skip_validation(self):
+        ms = tskit.MetadataSchema({"codec": "json+binary", "properties": {"foo": {}}})
+        assert not ms._bypass_validation
+        with patch.object(ms, "_validate_row", return_value=True) as mocked_validate:
+            ms.validate_and_encode_row({"_binary": b""})
+            assert mocked_validate.call_count == 1
+
+    def test_binary_requires_buffer_protocol(self):
+        ms = tskit.MetadataSchema({"codec": "json+binary"})
+        with pytest.raises(
+            exceptions.MetadataEncodingError,
+            match="_binary must be bytes-like \\(buffer protocol\\)",
+        ):
+            ms.validate_and_encode_row({"_binary": "not bytes"})
+
+    def test_decode_version_mismatch(self):
+        ms = tskit.MetadataSchema({"codec": "json+binary"})
+        header = metadata.JSONBinaryCodec._HDR.pack(
+            metadata.JSONBinaryCodec.MAGIC,
+            metadata.JSONBinaryCodec.VERSION + 1,
+            len(b"{}"),
+            0,
+        )
+        with pytest.raises(
+            ValueError,
+            match="Unsupported json\\+binary version",
+        ):
+            ms.decode_row(header + b"{}")
+
+
 class TestStructCodec:
     def encode_decode(self, method_name, sub_schema, obj, buffer):
         assert (

python/tskit/metadata.py

@@ -195,6 +195,72 @@ def decode(self, data: bytes) -> bytes:
         return data
 
 
+class JSONBinaryCodec(JSONCodec):
+    """
+    A JSON codec that optionally packs a single binary blob alongside the
+    canonical JSON bytes. The JSON portion is validated using the normal JSON
+    schema rules; a reserved top-level key "_binary" is ignored for validation
+    purposes and, if present at encode time, is stored as raw bytes appended
+    after a small header and the JSON payload.
+
+    On encode, callers MUST supply a top-level "_binary" bytes-like value,
+    even if zero length. On decode, the returned object will include a
+    "_binary" key whose value is a memoryview over the decoded bytes.
+    """
+
+    MAGIC = b"JBLB"
+    VERSION = 1
+    _HDR = struct.Struct("<4sBQQ")  # magic, version, json_len, blob_len
+
+    # Use the same validator behavior as JSONCodec; we do not special-case
+    # validation for the reserved _binary key. If users set additionalProperties
+    # to False, providing _binary will fail validation unless declared.
+
+    def encode(self, obj: Any) -> bytes:
+        # Require a top-level _binary bytes-like entry; zero-length allowed
+        if not isinstance(obj, dict) or "_binary" not in obj:
+            raise exceptions.MetadataEncodingError(
+                "json+binary requires top-level '_binary' bytes-like value"
+            )
+        try:
+            blob_bytes = memoryview(obj["_binary"]).tobytes()
+        except TypeError as e:
+            raise exceptions.MetadataEncodingError(
+                "_binary must be bytes-like (buffer protocol)"
+            ) from e
+
+        try:
+            json_bytes = tskit.canonical_json(
+                {k: v for k, v in obj.items() if k != "_binary"}
+            ).encode()
+        except TypeError as e:
+            raise exceptions.MetadataEncodingError(
+                f"Could not encode metadata of type {str(e).split()[3]}"
+            )
+
+        header = self._HDR.pack(
+            self.MAGIC, self.VERSION, len(json_bytes), len(blob_bytes)
+        )
+        return header + json_bytes + blob_bytes
+
+    def decode(self, encoded: bytes) -> Any:
+        if len(encoded) >= self._HDR.size and encoded[:4] == self.MAGIC:
+            _, version, jlen, blen = self._HDR.unpack_from(encoded)
+            if version != self.VERSION:
+                raise ValueError("Unsupported json+binary version")
+            start = self._HDR.size
+            json_bytes = encoded[start : start + jlen]
+            blob_bytes = encoded[start + jlen : start + jlen + blen]
+
+            result = super().decode(json_bytes)
+            result["_binary"] = memoryview(blob_bytes)
+            return result
+        raise ValueError("Invalid json+binary payload: missing magic header")
+
+
+register_metadata_codec(JSONBinaryCodec, "json+binary")
+
+
 def binary_format_validator(validator, types, instance, schema):
     # We're hooking into jsonschemas validation code here, which works by creating
     # generators of exceptions, hence the yielding