feat(hugr-py): Define typed Metadata protocol

hugr-core/src/metadata.rs

@@ -22,6 +22,9 @@
 //! let payload = hugr.get_metadata::<SomeMetadata>(hugr.module_root());
 //! assert_eq!(payload, Some("payload"));
 //! ```
+//
+// When adding new metadata keys, they should be re-exported by the python bindings.
+// See hugr-py/rust/metadata.rs
 
 /// Arbitrary metadata entry for a node.
 ///

hugr-py/Cargo.toml

@@ -22,6 +22,9 @@ bench = false
 [dependencies]
 bumpalo = { workspace = true, features = ["collections"] }
 hugr-cli = { version = "0.24.3", path = "../hugr-cli", default-features = false }
-hugr-model = { version = "0.24.3", path = "../hugr-model", features = ["pyo3"] }
+hugr-model = { version = "0.24.3", path = "../hugr-model", default-features = false, features = [
+    "pyo3",
+] }
+hugr-core = { version = "0.24.3", path = "../hugr-core", default-features = false }
 pastey.workspace = true
 pyo3 = { workspace = true, features = ["extension-module", "abi3-py310"] }

hugr-py/rust/lib.rs

@@ -1,12 +1,14 @@
 //! Supporting Rust library for the hugr Python bindings.
 
+mod metadata;
 mod model;
 
 use pyo3::pymodule;
 
 #[pymodule]
 mod _hugr {
-
+    #[pymodule_export]
+    use super::metadata::metadata;
     #[pymodule_export]
     use super::model::model;
 }

hugr-py/rust/metadata.rs

@@ -0,0 +1,25 @@
+//! Bindings for metadata keys defined in the hugr-core crate.
+
+#[pyo3::pymodule(submodule)]
+#[pyo3(module = "hugr._hugr.metadata")]
+pub mod metadata {
+    use hugr_core::metadata::Metadata;
+    use pyo3::types::{PyAnyMethods, PyModule};
+    use pyo3::{Bound, PyResult, Python};
+
+    /// Hack: workaround for <https://github.com/PyO3/pyo3/issues/759>
+    #[pymodule_init]
+    fn init(m: &Bound<'_, PyModule>) -> PyResult<()> {
+        Python::attach(|py| {
+            py.import("sys")?
+                .getattr("modules")?
+                .set_item("hugr._hugr.metadata", m)
+        })
+    }
+
+    #[pymodule_export]
+    const HUGR_GENERATOR: &str = hugr_core::metadata::HugrGenerator::KEY;
+
+    #[pymodule_export]
+    const HUGR_USED_EXTENSIONS: &str = hugr_core::metadata::HugrUsedExtensions::KEY;
+}

hugr-py/src/hugr/_hugr/metadata.pyi

@@ -0,0 +1,2 @@
+HUGR_GENERATOR: str
+HUGR_USED_EXTENSIONS: str

hugr-py/src/hugr/build/dfg.py

@@ -23,6 +23,7 @@
 
     from hugr.build.function import Module
     from hugr.hugr.node_port import Node, OutPort, PortOffset, ToNode, Wire
+    from hugr.metadata import NodeMetadata
     from hugr.tys import TypeParam, TypeRow
 
     from .cfg import Cfg
@@ -195,7 +196,11 @@ def inputs(self) -> list[OutPort]:
         return [self.input_node.out(i) for i in range(len(self._input_op().types))]
 
     def add_op(
-        self, op: ops.DataflowOp, /, *args: Wire, metadata: dict[str, Any] | None = None
+        self,
+        op: ops.DataflowOp,
+        /,
+        *args: Wire,
+        metadata: dict[str, Any] | NodeMetadata | None = None,
     ) -> Node:
         """Add a dataflow operation to the graph, wiring in input ports.
 

hugr-py/src/hugr/build/function.py

@@ -11,6 +11,7 @@
 
 if TYPE_CHECKING:
     from hugr.hugr.node_port import Node
+    from hugr.metadata import NodeMetadata
     from hugr.tys import PolyFuncType, Type, TypeBound, TypeParam, TypeRow
 
 __all__ = ["Function", "Module"]
@@ -91,6 +92,6 @@ def add_alias_decl(self, name: str, bound: TypeBound) -> Node:
         return self.hugr.add_node(ops.AliasDecl(name, bound), self.hugr.module_root)
 
     @property
-    def metadata(self) -> dict[str, object]:
+    def metadata(self) -> NodeMetadata:
         """Metadata associated with this module."""
         return self.hugr.entrypoint.metadata

hugr-py/src/hugr/envelope.py

@@ -35,9 +35,10 @@
 import json
 from dataclasses import dataclass
 from enum import Enum
-from typing import TYPE_CHECKING, ClassVar
+from typing import TYPE_CHECKING, Any, ClassVar
 
 import pyzstd
+from semver import Version
 
 import hugr._hugr.model as rust
 
@@ -263,3 +264,97 @@ def _make_header(self) -> EnvelopeHeader:
 # These can only be initialized _after_ the class is defined.
 EnvelopeConfig.TEXT = EnvelopeConfig(format=EnvelopeFormat.JSON, zstd=None)
 EnvelopeConfig.BINARY = EnvelopeConfig(format=EnvelopeFormat.MODEL_WITH_EXTS, zstd=0)
+
+
+@dataclass(frozen=True)
+class GeneratorDesc:
+    """Description of the generator that defined the HUGR module.
+
+    These are stored at the module root node metadata under the
+    :class:`hugr.metadata.HugrGenerator` entry.
+    """
+
+    name: str
+    version: Version | None
+
+    def _to_json(self) -> dict[str, str]:
+        """Encodes the generator as a dictionary of native types that can be
+        serialized by `json.dump`.
+        """
+        if self.version is None:
+            return {
+                "name": self.name,
+            }
+        else:
+            return {
+                "name": self.name,
+                "version": str(self.version),
+            }
+
+    @classmethod
+    def _from_json(cls, value: Any) -> GeneratorDesc:
+        """Decodes the generator from a native types obtained from `json.load`."""
+        if isinstance(value, str):
+            return GeneratorDesc(name=value, version=None)
+
+        if not isinstance(value, dict):
+            msg = (
+                "Expected generator metadata to be a string or a dict,"
+                + " but got {type(value)}"
+            )
+            raise TypeError(msg)
+
+        fallback_name = " ".join(f"{k}: {v}" for k, v in value.items())
+        if "name" not in value or any(k != "name" and k != "version" for k in value):
+            return GeneratorDesc(name=fallback_name, version=None)
+        if "version" in value:
+            try:
+                version = Version.parse(value["version"])
+            except ValueError:
+                return GeneratorDesc(name=fallback_name, version=None)
+            return GeneratorDesc(name=value["name"], version=version)
+        else:
+            return GeneratorDesc(name=value["name"], version=None)
+
+
+@dataclass
+class ExtensionDesc:
+    """High level description of a HUGR extension.
+
+    A list of these is stored at the module root node metadata under the
+    :class:`hugr.metadata.HugrUsedExtensions` entry.
+    """
+
+    name: str
+    version: Version
+
+    def _to_json(self) -> dict[str, str]:
+        """Encodes the extension as a dictionary of native types that can be
+        serialized by `json.dump`.
+        """
+        return {
+            "name": self.name,
+            "version": str(self.version),
+        }
+
+    @classmethod
+    def _from_json(cls, value: Any) -> ExtensionDesc:
+        """Decodes the extension from a native types obtained from `json.load`."""
+        if not isinstance(value, dict):
+            msg = f"Expected extension metadata to be a dict, but got {type(value)}"
+            raise TypeError(msg)
+        if "name" not in value:
+            msg = (
+                "Expected extension metadata to be a dict with a 'name' key,"
+                + f" but got {value}"
+            )
+            raise TypeError(msg)
+        if "version" not in value:
+            msg = (
+                "Expected extension metadata to be a dict with a 'version' key,"
+                + f" but got {value}"
+            )
+            raise TypeError(msg)
+        return ExtensionDesc(
+            name=value["name"], version=Version.parse(value["version"])
+        )

hugr-py/src/hugr/hugr/base.py

@@ -2,6 +2,7 @@
 
 from __future__ import annotations
 
+import copy
 import json
 from collections.abc import Iterable, Iterator, Mapping
 from dataclasses import dataclass, field
@@ -29,6 +30,7 @@
     read_envelope_hugr_str,
 )
 from hugr.exceptions import ParentBeforeChild
+from hugr.metadata import NodeMetadata
 from hugr.ops import (
     CFG,
     Call,
@@ -79,7 +81,7 @@ class NodeData:
     _num_inps: int = field(default=0, repr=False)
     _num_outs: int = field(default=0, repr=False)
     children: list[Node] = field(default_factory=list, repr=False)
-    metadata: dict[str, Any] = field(default_factory=dict)
+    metadata: NodeMetadata = field(default_factory=NodeMetadata)
 
     def _to_serial(self, node: Node) -> SerialOp:
         o = self.op._to_serial(self.parent if self.parent else node)
@@ -344,16 +346,20 @@ def _add_node(
         op: Op,
         parent: ToNode | None = None,
         num_outs: int | None = None,
-        metadata: dict[str, Any] | None = None,
+        metadata: dict[str, Any] | NodeMetadata | None = None,
     ) -> Node:
         parent = parent.to_node() if parent else None
-        node_data = NodeData(op, parent, metadata=metadata or {})
+
+        if metadata is None or isinstance(metadata, dict):
+            metadata = NodeMetadata(metadata)
+
+        node_data = NodeData(op, parent, metadata=metadata)
 
         if self._free_nodes:
             node = self._free_nodes.pop()
             self._nodes[node.idx] = node_data
         else:
-            node = Node(len(self._nodes), {})
+            node = Node(len(self._nodes), NodeMetadata())
             self._nodes.append(node_data)
         node._num_out_ports = num_outs
         node._metadata = node_data.metadata
@@ -402,7 +408,7 @@ def add_node(
         op: Op,
         parent: ToNode | None = None,
         num_outs: int | None = None,
-        metadata: dict[str, Any] | None = None,
+        metadata: dict[str, Any] | NodeMetadata | None = None,
     ) -> Node:
         """Add a node to the HUGR.
 
@@ -423,7 +429,7 @@ def add_const(
         self,
         value: Value,
         parent: ToNode | None = None,
-        metadata: dict[str, Any] | None = None,
+        metadata: dict[str, Any] | NodeMetadata | None = None,
     ) -> Node:
         """Add a constant node to the HUGR.
 
@@ -474,7 +480,7 @@ def delete_node(self, node: ToNode) -> NodeData | None:
         weight, self._nodes[node.idx] = self._nodes[node.idx], None
 
         # Free up the metadata dictionary
-        node._metadata = {}
+        node._metadata = NodeMetadata()
 
         self._free_nodes.append(node)
         return weight
@@ -921,7 +927,7 @@ def insert_hugr(self, hugr: Hugr, parent: ToNode | None = None) -> dict[Node, No
                 node_data.op,
                 node_parent,
                 num_outs=node_data._num_outs,
-                metadata=node_data.metadata,
+                metadata=copy.copy(node_data.metadata),
             )
 
         for src, dst in hugr._links.items():
@@ -964,8 +970,8 @@ def _serialize_link(
             serial_idx = len(nodes)
 
             # non contiguous indices will be erased
-            nodes.append(data._to_serial(Node(serial_idx, {})))
-            metadata.append(data.metadata if data.metadata else None)
+            nodes.append(data._to_serial(Node(serial_idx, NodeMetadata())))
+            metadata.append(data.metadata.as_dict() if data.metadata else None)
             if self.entrypoint == node:
                 entrypoint = serial_idx
 
@@ -1022,12 +1028,11 @@ def _from_serial(cls, serial: SerialHugr) -> Hugr:
         """Load a HUGR from a serialized form."""
         assert serial.nodes, "The encoded Hugr is empty"
 
-        def get_meta(idx: int) -> dict[str, Any]:
-            if not serial.metadata:
-                return {}
-            if idx < len(serial.metadata):
-                return serial.metadata[idx] or {}
-            return {}
+        def get_meta(idx: int) -> NodeMetadata:
+            if serial.metadata and idx < len(serial.metadata):
+                return NodeMetadata(serial.metadata[idx] or {})
+            else:
+                return NodeMetadata()
 
         # The first node is always the HUGR root.
         root_node = serial.nodes[0]

hugr-py/src/hugr/hugr/node_port.py

@@ -15,6 +15,8 @@
 
 from typing_extensions import Self
 
+from hugr.metadata import NodeMetadata
+
 if TYPE_CHECKING:
     from collections.abc import Iterator
 
@@ -142,7 +144,7 @@ def port(self, offset: PortOffset, direction: Direction) -> InPort | OutPort:
             return self.out(offset)
 
     @property
-    def metadata(self) -> dict[str, object]:
+    def metadata(self) -> NodeMetadata:
         """Metadata associated with this node."""
         return self.to_node()._metadata
 
@@ -154,8 +156,8 @@ class Node(ToNode):
     """
 
     idx: NodeIdx
-    _metadata: dict[str, object] = field(
-        repr=False, compare=False, default_factory=dict
+    _metadata: NodeMetadata = field(
+        repr=False, compare=False, default_factory=NodeMetadata
     )
     _num_out_ports: int | None = field(default=None, compare=False, repr=False)