feat: Add serialization schema for metadata (#1038)

fixes #954. We support deserializing HUGRs that do not have a `metadata` field. We could optimise a little by not serializing a `metadata` field when all node metadata fields are null, but for now we don't bother. --------- Co-authored-by: Agustín Borgna <[email protected]>
CQCL · May 14, 2024 · 19bac62 · 19bac62
1 parent 10e3158
commit 19bac62
Show file tree

Hide file tree

Showing 4 changed files with 62 additions and 18 deletions.
diff --git a/hugr-py/src/hugr/serialization/serial_hugr.py b/hugr-py/src/hugr/serialization/serial_hugr.py
@@ -1,21 +1,22 @@
 from typing import Any, Literal
 
-from pydantic import BaseModel, Field, ConfigDict
+from pydantic import Field, ConfigDict
 
 from .ops import NodeID, OpType, classes as ops_classes
-from .tys import model_rebuild
+from .tys import model_rebuild, ConfiguredBaseModel
 import hugr
 
 Port = tuple[NodeID, int | None]  # (node, offset)
 Edge = tuple[Port, Port]
 
 
-class SerialHugr(BaseModel):
+class SerialHugr(ConfiguredBaseModel):
     """A serializable representation of a Hugr."""
 
     version: Literal["v1"] = "v1"
     nodes: list[OpType]
     edges: list[Edge]
+    metadata: list[dict[str, Any] | None] | None = None
     encoder: str | None = Field(
         default=None, description="The name of the encoder used to generate the Hugr."
     )

diff --git a/hugr/src/hugr/serialize.rs b/hugr/src/hugr/serialize.rs
@@ -14,7 +14,7 @@ use portgraph::{Direction, LinkError, PortView};
 
 use serde::{Deserialize, Deserializer, Serialize};
 
-use super::{HugrMut, HugrView};
+use super::{HugrMut, HugrView, NodeMetadataMap};
 
 /// A wrapper over the available HUGR serialization formats.
 ///
@@ -61,11 +61,8 @@ struct SerHugrV1 {
     /// for each edge: (src, src_offset, tgt, tgt_offset)
     edges: Vec<[(Node, Option<u16>); 2]>,
     /// for each node: (metadata)
-    //
-    // TODO: Update to Vec<Option<Map<String,Value>>> to more closely
-    // match the internal representation.
     #[serde(default)]
-    metadata: Vec<serde_json::Value>,
+    metadata: Option<Vec<Option<NodeMetadataMap>>>,
     /// A metadata field with the package identifier that encoded the HUGR.
     #[serde(default)]
     encoder: Option<String>,
@@ -173,7 +170,7 @@ impl TryFrom<&Hugr> for SerHugrV1 {
         }
 
         let mut nodes = vec![None; hugr.node_count()];
-        let mut metadata = vec![serde_json::Value::Null; hugr.node_count()];
+        let mut metadata = vec![None; hugr.node_count()];
         for n in hugr.nodes() {
             let parent = node_rekey[&hugr.get_parent(n).unwrap_or(n)];
             let opt = hugr.get_nodetype(n);
@@ -183,11 +180,7 @@ impl TryFrom<&Hugr> for SerHugrV1 {
                 input_extensions: opt.input_extensions.clone(),
                 op: opt.op.clone(),
             });
-            let node_metadata = hugr.metadata.get(n.pg_index()).clone();
-            metadata[new_node] = match node_metadata {
-                Some(m) => serde_json::Value::Object(m.clone()),
-                None => serde_json::Value::Null,
-            };
+            metadata[new_node].clone_from(hugr.metadata.get(n.pg_index()));
         }
         let nodes = nodes
             .into_iter()
@@ -222,7 +215,7 @@ impl TryFrom<&Hugr> for SerHugrV1 {
         Ok(Self {
             nodes,
             edges,
-            metadata,
+            metadata: Some(metadata),
             encoder,
         })
     }
@@ -263,9 +256,13 @@ impl TryFrom<SerHugrV1> for Hugr {
             );
         }
 
-        for (node, metadata) in metadata.into_iter().enumerate() {
-            let node = portgraph::NodeIndex::new(node);
-            hugr.metadata[node] = metadata.as_object().cloned();
+        if let Some(metadata) = metadata {
+            for (node, metadata) in metadata.into_iter().enumerate() {
+                if let Some(metadata) = metadata {
+                    let node = portgraph::NodeIndex::new(node);
+                    hugr.metadata[node] = Some(metadata);
+                }
+            }
         }
 
         let unwrap_offset = |node: Node, offset, dir, hugr: &Hugr| -> Result<usize, Self::Error> {

diff --git a/specification/schema/hugr_schema_strict_v1.json b/specification/schema/hugr_schema_strict_v1.json
@@ -2259,6 +2259,7 @@
             "type": "object"
         }
     },
+    "additionalProperties": false,
     "description": "A serializable representation of a Hugr.",
     "properties": {
         "version": {
@@ -2328,6 +2329,28 @@
             "title": "Edges",
             "type": "array"
         },
+        "metadata": {
+            "anyOf": [
+                {
+                    "items": {
+                        "anyOf": [
+                            {
+                                "type": "object"
+                            },
+                            {
+                                "type": "null"
+                            }
+                        ]
+                    },
+                    "type": "array"
+                },
+                {
+                    "type": "null"
+                }
+            ],
+            "default": null,
+            "title": "Metadata"
+        },
         "encoder": {
             "anyOf": [
                 {

diff --git a/specification/schema/hugr_schema_v1.json b/specification/schema/hugr_schema_v1.json
@@ -2259,6 +2259,7 @@
             "type": "object"
         }
     },
+    "additionalProperties": true,
     "description": "A serializable representation of a Hugr.",
     "properties": {
         "version": {
@@ -2328,6 +2329,28 @@
             "title": "Edges",
             "type": "array"
         },
+        "metadata": {
+            "anyOf": [
+                {
+                    "items": {
+                        "anyOf": [
+                            {
+                                "type": "object"
+                            },
+                            {
+                                "type": "null"
+                            }
+                        ]
+                    },
+                    "type": "array"
+                },
+                {
+                    "type": "null"
+                }
+            ],
+            "default": null,
+            "title": "Metadata"
+        },
         "encoder": {
             "anyOf": [
                 {