feat: Add python definitions for the TKET metadata keys

tket-py/src/lib.rs

@@ -1,4 +1,5 @@
 //! Python bindings for TKET.
+pub mod metadata;
 pub mod ops;
 pub mod optimiser;
 pub mod passes;
@@ -13,6 +14,7 @@ use pyo3::prelude::*;
 /// The Python bindings to TKET.
 #[pymodule]
 fn _tket(py: Python, m: &Bound<PyModule>) -> PyResult<()> {
+    add_submodule(py, m, metadata::module(py)?)?;
     add_submodule(py, m, state::module(py)?)?;
     add_submodule(py, m, ops::module(py)?)?;
     add_submodule(py, m, optimiser::module(py)?)?;

tket-py/src/metadata.rs

@@ -0,0 +1,22 @@
+//! Bindings for metadata keys defined in the `tket` crate.
+
+use hugr::metadata::Metadata;
+use pyo3::prelude::*;
+use tket::metadata::{
+    BitRegisters, CircuitRewriteTraces, InputParameters, MaxQubits, OpGroup, Phase, QubitRegisters,
+    Unitary,
+};
+
+/// The module definition.
+pub fn module(py: Python<'_>) -> PyResult<Bound<'_, PyModule>> {
+    let m = PyModule::new(py, "metadata")?;
+    m.add("MAX_QUBITS", MaxQubits::KEY)?;
+    m.add("CIRCUIT_REWRITE_TRACES", CircuitRewriteTraces::KEY)?;
+    m.add("UNITARY", Unitary::KEY)?;
+    m.add("INPUT_PARAMETERS", InputParameters::KEY)?;
+    m.add("OP_GROUP", OpGroup::KEY)?;
+    m.add("BIT_REGISTERS", BitRegisters::KEY)?;
+    m.add("QUBIT_REGISTERS", QubitRegisters::KEY)?;
+    m.add("PHASE", Phase::KEY)?;
+    Ok(m)
+}

tket-py/tket/__init__.py

@@ -9,10 +9,11 @@
 [crates.io](https://crates.io/crates/tket).
 """
 
-from . import passes
+from . import passes, metadata
 
 __all__ = [
     "passes",
+    "metadata",
 ]
 
 

tket-py/tket/_tket/metadata.pyi

@@ -0,0 +1,8 @@
+MAX_QUBITS: str
+CIRCUIT_REWRITE_TRACES: str
+UNITARY: str
+INPUT_PARAMETERS: str
+OP_GROUP: str
+BIT_REGISTERS: str
+QUBIT_REGISTERS: str
+PHASE: str

tket-py/tket/metadata.py

@@ -0,0 +1,151 @@
+"""Metadata values defined by the TKET compiler.
+
+Examples:
+    >>> from hugr import Hugr
+    >>> from tket.metadata import InputParameters, MaxQubits, Phase, QubitRegisters
+    >>>
+    >>> hugr = Hugr()
+    >>> node = hugr[hugr.module_root]
+    >>>
+    >>> node.metadata[MaxQubits] = 3
+    >>> node.metadata[InputParameters] = ["theta", "phi"]
+    >>> node.metadata[QubitRegisters] = [("q", [0]), ("ancilla", [1])]
+    >>> node.metadata[Phase] = "1/2"
+    >>> node.metadata[MaxQubits]
+    3
+    >>> node.metadata.get(QubitRegisters)
+    [('q', [0]), ('ancilla', [1])]
+"""
+# Changes to this file **MUST** be reflected in `tket/src/metadata.rs`
+
+from __future__ import annotations
+
+from typing import TYPE_CHECKING, TypeAlias, TypedDict
+
+from hugr.metadata import Metadata
+from ._tket import metadata as _metadata
+
+if TYPE_CHECKING:
+    from hugr.utils import JsonType
+
+
+__all__ = [
+    "RewriteTraceValue",
+    "MaxQubits",
+    "CircuitRewriteTraces",
+    "Unitary",
+    "InputParameters",
+    "OpGroup",
+    "BitRegisters",
+    "QubitRegisters",
+    "Phase",
+]
+
+
+# Identifier for a TKET1 qubit register element.
+#
+# This can be passed to `pytket.unit_id.Qubit.from_list`
+PytketQubit: TypeAlias = tuple[str, list[int]]
+# Identifier for a TKET1 bit register element.
+#
+# This can be passed to `pytket.unit_id.Bit.from_list`
+PytketBit: TypeAlias = tuple[str, list[int]]
+
+
+class RewriteTraceValue(TypedDict):
+    """Serialized rewrite trace metadata entry."""
+
+    individual_matches: int
+
+
+class MaxQubits(Metadata[int]):
+    """Metadata key for the number of qubits required to execute a HUGR node."""
+
+    KEY = _metadata.MAX_QUBITS
+
+
+class CircuitRewriteTraces(Metadata[list[RewriteTraceValue]]):
+    """Metadata key for rewrite traces recorded during circuit transformation."""
+
+    KEY = _metadata.CIRCUIT_REWRITE_TRACES
+
+
+class Unitary(Metadata[int]):
+    """Metadata key for unitary/modifier flags stored on a HUGR node."""
+
+    KEY = _metadata.UNITARY
+
+
+class InputParameters(Metadata[list[str]]):
+    """Metadata key for explicit names of input parameter wires."""
+
+    KEY = _metadata.INPUT_PARAMETERS
+
+
+class OpGroup(Metadata[str]):
+    """Metadata key for the pytket ``opgroup`` field on a decoded operation."""
+
+    KEY = _metadata.OP_GROUP
+
+
+class BitRegisters(Metadata[list[PytketBit]]):
+    """Metadata key for explicit names of input bit registers."""
+
+    KEY = _metadata.BIT_REGISTERS
+
+    @classmethod
+    def to_json(cls, value: list[PytketBit]) -> JsonType:
+        return _store_pytket_register(value)
+
+    @classmethod
+    def from_json(cls, value: JsonType) -> list[PytketBit]:
+        return _read_pytket_register(cls.KEY, value)
+
+
+class QubitRegisters(Metadata[list[PytketQubit]]):
+    """Metadata key for explicit names of input qubit registers."""
+
+    KEY = _metadata.QUBIT_REGISTERS
+
+    @classmethod
+    def to_json(cls, value: list[PytketQubit]) -> JsonType:
+        return _store_pytket_register(value)
+
+    @classmethod
+    def from_json(cls, value: JsonType) -> list[PytketQubit]:
+        return _read_pytket_register(cls.KEY, value)
+
+
+class Phase(Metadata[str]):
+    """Metadata key for the serialized TKET1 global phase expression."""
+
+    KEY = _metadata.PHASE
+
+
+def _store_pytket_register(value: list[tuple[str, list[int]]]) -> JsonType:
+    return [[name, indices] for name, indices in value]  # type: ignore
+
+
+def _read_pytket_register(key: str, value: JsonType) -> list[tuple[str, list[int]]]:
+    if not isinstance(value, list):
+        raise TypeError(f"Expected {key} metadata to be a list, but got {type(value)}")
+
+    registers: list[tuple[str, list[int]]] = []
+    for entry in value:
+        if not isinstance(entry, list) or len(entry) != 2:
+            raise TypeError(
+                f"Expected each {key} metadata entry to be [name, [indices...]], but got {entry!r}"
+            )
+        name, indices = entry
+        if not isinstance(name, str):
+            raise TypeError(
+                f"Expected {key} register name to be a string, but got {type(name)}"
+            )
+        if not isinstance(indices, list) or not all(
+            isinstance(index, int) for index in indices
+        ):
+            raise TypeError(
+                f"Expected {key} register indices to be a list of integers, but got {indices!r}"
+            )
+        registers.append((name, indices))  # type: ignore
+    return registers

tket/src/metadata.rs

@@ -1,4 +1,32 @@
 //! Collection of metadata keys used throughout tket.
+//!
+//! # Example
+//!
+//! ```rust
+//! use tket::metadata;
+//! use hugr::{Hugr, HugrView};
+//! # use hugr::hugr::hugrmut::HugrMut;
+//! # use hugr::types::Signature;
+//! # use tket_json_rs::register::{ElementId, Qubit};
+//!
+//! let mut hugr = Hugr::new();
+//! let node = hugr.entrypoint();
+//!
+//! hugr.set_metadata::<metadata::MaxQubits>(node, 3);
+//! hugr.set_metadata::<metadata::InputParameters>(node, vec!["theta".to_string()]);
+//! hugr.set_metadata::<metadata::QubitRegisters>(
+//!     node,
+//!     vec![Qubit::from(ElementId("q".to_string(), vec![0]))],
+//! );
+//!
+//! assert_eq!(hugr.get_metadata::<metadata::MaxQubits>(node), Some(3));
+//! assert_eq!(
+//!     hugr.get_metadata::<metadata::InputParameters>(node),
+//!     Some(vec!["theta".to_string()]),
+//! );
+//! ```
+//
+// Changes to this file **MUST** be reflected in `tket-py/tket/metadata.py`
 
 use crate::rewrite::trace::RewriteTrace;
 use hugr_core::metadata::Metadata;