feat: Export metadata keys from hugr-core

Author: aborgna-q

Cargo.lock

@@ -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-core/src/metadata.rs

@@ -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/Cargo.toml

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

hugr-py/rust/lib.rs

@@ -0,0 +1,26 @@
+//! Bindings for metadata keys defined in the hugr-core crate.
+
+#[pyo3::pymodule]
+#[pyo3(submodule)]
+#[pyo3(module = "hugr._hugr.metadata")]
+pub mod metadata {
+    use hugr_core::metadata::Metadata;
+    use pyo3::types::{PyAnyMethods, PyModule};
+    use pyo3::{Bound, PyResult, Python};
+
+    #[pymodule_export]
+    const HUGR_GENERATOR: &str = hugr_core::metadata::HugrGenerator::KEY;
+
+    #[pymodule_export]
+    const HUGR_USED_EXTENSIONS: &str = hugr_core::metadata::HugrUsedExtensions::KEY;
+
+    /// 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)
+        })
+    }
+}

hugr-py/rust/metadata.rs

@@ -5,7 +5,8 @@ use hugr_cli::RunWithIoError;
 use pyo3::exceptions::{PyException, PyValueError};
 use pyo3::{PyErr, PyResult, create_exception, pymodule};
 
-#[pymodule(submodule)]
+#[pymodule]
+#[pyo3(submodule)]
 pub mod model {
     use hugr_cli::CliArgs;
     use hugr_model::v0::ast;

hugr-py/rust/model.rs

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

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

@@ -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/envelope.py

@@ -5,18 +5,40 @@
 from dataclasses import dataclass
 from typing import TYPE_CHECKING, Any, ClassVar, Protocol, TypeVar, overload
 
+import hugr._hugr.metadata as rust_metadata
+from hugr.envelope import ExtensionDesc, GeneratorDesc
+
 if TYPE_CHECKING:
     from collections.abc import Iterable, Iterator
 
-MetaCovariant = TypeVar("MetaCovariant", covariant=True)
 Meta = TypeVar("Meta")
 
 
-class Metadata(Protocol[MetaCovariant]):
-    """Metadata for a HUGR node."""
+class Metadata(Protocol[Meta]):
+    """Metadata for a HUGR node.
+
+    This is a protocol for metadata entries that defines a unique key to
+    identify the entry, and the type of the value.
+
+    Values in a hugr are encoded using json. When the value type is not a
+    primitive type, `to_json` and `from_json` must be implemented to serialize
+    and deserialize the value.
+
+    Args:
+        value: The value of the metadata.
+    """
 
     KEY: ClassVar[str]
-    TYPE: ClassVar[type]
+
+    @classmethod
+    def to_json(cls, value: Meta) -> Any:
+        """Serialize the metadata value to a json value."""
+        return value
+
+    @classmethod
+    def from_json(cls, value: Any) -> Meta:
+        """Deserialize the metadata value from the stored json value."""
+        return value
 
 
 @dataclass
@@ -35,11 +57,19 @@ def __init__(self, metadata: dict[str, Any] | None = None) -> None:
     @overload
     def get(self, key: str, default: Any | None = None) -> Any | None: ...
     @overload
-    def get(self, key: Metadata[Meta], default: Meta | None = None) -> Meta | None: ...
-    def get(self, key: str | Metadata[Meta], default: Any | None = None) -> Meta | None:
-        if not isinstance(key, str):
-            key = key.KEY
-        return self._dict.get(key, default)
+    def get(
+        self, key: type[Metadata[Meta]], default: Meta | None = None
+    ) -> Meta | None: ...
+    def get(
+        self, key: str | type[Metadata[Meta]], default: Any | None = None
+    ) -> Any | None:
+        if isinstance(key, str):
+            return self._dict.get(key, default)
+        elif key.KEY in self._dict:
+            val = self._dict[key.KEY]
+            return key.from_json(val)
+        else:
+            return None
 
     def items(self) -> Iterable[tuple[str, Any]]:
         return self._dict.items()
@@ -50,34 +80,81 @@ def as_dict(self) -> dict[str, Any]:
     @overload
     def __getitem__(self, key: str) -> Any: ...
     @overload
-    def __getitem__(self, key: Metadata[Meta]) -> Meta: ...
-    def __getitem__(self, key: str | Metadata[Meta]) -> Any:
-        if not isinstance(key, str):
-            key = key.KEY
-        return self._dict[key]
+    def __getitem__(self, key: type[Metadata[Meta]]) -> Meta: ...
+    def __getitem__(self, key: str | type[Metadata[Meta]]) -> Any:
+        if isinstance(key, str):
+            return self._dict[key]
+        else:
+            val = self._dict[key.KEY]
+            return key.from_json(val)
 
     @overload
     def __setitem__(self, key: str, value: Any) -> None: ...
     @overload
-    def __setitem__(self, key: Metadata[Meta], value: Meta) -> None: ...
-    def __setitem__(self, key: str | Metadata[Meta], value: Any) -> None:
-        if not isinstance(key, str):
-            if not isinstance(value, key.TYPE):
-                error = f"Value for metadata key {key.KEY} must be of type {key.TYPE}"
-                raise TypeError(error)
-            key = key.KEY
-        self._dict[key] = value
+    def __setitem__(self, key: type[Metadata[Meta]], value: Meta) -> None: ...
+    def __setitem__(self, key: str | type[Metadata[Meta]], value: Any) -> None:
+        if isinstance(key, str):
+            self._dict[key] = value
+        else:
+            json_value = key.to_json(value)
+            self._dict[key.KEY] = json_value
 
     def __iter__(self) -> Iterator[str]:
         return iter(self._dict)
 
     def __len__(self) -> int:
         return len(self._dict)
 
-    def __contains__(self, key: str | Metadata[Meta]) -> bool:
+    def __contains__(self, key: str | type[Metadata[Meta]]) -> bool:
         if not isinstance(key, str):
             key = key.KEY
         return key in self._dict
 
     def __repr__(self) -> str:
         return f"NodeMetadata({self._dict})"
+
+
+# --- Core metadata keys ---
+
+
+class HugrGenerator(Metadata[GeneratorDesc]):
+    """Metadata describing the generator that defined the HUGR module.
+
+    This value is only valid when set at the module root node.
+    """
+
+    KEY = rust_metadata.HUGR_GENERATOR
+
+    @classmethod
+    def to_json(cls, value: GeneratorDesc) -> dict[str, str]:
+        return value._to_json()
+
+    @classmethod
+    def from_json(cls, value: Any) -> GeneratorDesc:
+        return GeneratorDesc._from_json(value)
+
+
+class HugrUsedExtensions(Metadata[list[ExtensionDesc]]):
+    """Metadata storing the list of extensions required to define the HUGR.
+
+    This list may contain additional extensions that are no longer present in
+    the Hugr.
+
+    This value is only valid when set at the module root node.
+    """
+
+    KEY = rust_metadata.HUGR_USED_EXTENSIONS
+
+    @classmethod
+    def to_json(cls, value: list[ExtensionDesc]) -> list[dict[str, str]]:
+        return [e._to_json() for e in value]
+
+    @classmethod
+    def from_json(cls, value: Any) -> list[ExtensionDesc]:
+        if not isinstance(value, list):
+            msg = (
+                "Expected UsedExtensions metadata to be a list,"
+                + f" but got {type(value)}"
+            )
+            raise TypeError(msg)
+        return [ExtensionDesc._from_json(e) for e in value]