feat: Allow importing from model data to python.

hugr-py/Cargo.toml

@@ -21,6 +21,7 @@ bench = false
 
 [dependencies]
 bumpalo = { workspace = true, features = ["collections"] }
+hugr-core = { version = "0.22.4", path = "../hugr-core", features = ["zstd"] }
 hugr-model = { version = "0.22.4", path = "../hugr-model", features = ["pyo3"] }
 pastey.workspace = true
 pyo3 = { workspace = true, features = ["extension-module", "abi3-py310"] }

hugr-py/rust/lib.rs

@@ -1,5 +1,9 @@
 //! Supporting Rust library for the hugr Python bindings.
 
+use hugr_core::{
+    envelope::{EnvelopeConfig, EnvelopeFormat, read_envelope, write_envelope},
+    std_extensions::STD_REG,
+};
 use hugr_model::v0::ast;
 use pyo3::{exceptions::PyValueError, prelude::*};
 
@@ -50,6 +54,17 @@ fn bytes_to_package(bytes: &[u8]) -> PyResult<ast::Package> {
     Ok(package)
 }
 
+/// Convert an envelope to a new envelope in JSON format.
+#[pyfunction]
+fn to_json_envelope(bytes: &[u8]) -> PyResult<Vec<u8>> {
+    let (_, pkg) =
+        read_envelope(bytes, &STD_REG).map_err(|err| PyValueError::new_err(err.to_string()))?;
+    let config_json = EnvelopeConfig::new(EnvelopeFormat::PackageJson);
+    let mut json_data: Vec<u8> = Vec::new();
+    write_envelope(&mut json_data, &pkg, config_json).unwrap();
+    Ok(json_data)
+}
+
 /// Returns the current version of the HUGR model format as a tuple of (major, minor, patch).
 #[pyfunction]
 fn current_model_version() -> (u64, u64, u64) {
@@ -79,5 +94,6 @@ fn _hugr(m: &Bound<'_, PyModule>) -> PyResult<()> {
     m.add_function(wrap_pyfunction!(symbol_to_string, m)?)?;
     m.add_function(wrap_pyfunction!(string_to_symbol, m)?)?;
     m.add_function(wrap_pyfunction!(current_model_version, m)?)?;
+    m.add_function(wrap_pyfunction!(to_json_envelope, m)?)?;
     Ok(())
 }

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

@@ -19,3 +19,4 @@ def string_to_package(string: str) -> hugr.model.Package: ...
 def package_to_bytes(package: hugr.model.Package) -> bytes: ...
 def bytes_to_package(binary: bytes) -> hugr.model.Package: ...
 def current_model_version() -> tuple[int, int, int]: ...
+def to_json_envelope(binary: bytes) -> bytes: ...

hugr-py/src/hugr/envelope.py

@@ -39,6 +39,8 @@
 
 import pyzstd
 
+import hugr._hugr as rust
+
 if TYPE_CHECKING:
     from hugr.hugr.base import Hugr
     from hugr.package import Package
@@ -112,8 +114,11 @@ def read_envelope(envelope: bytes) -> Package:
         case EnvelopeFormat.JSON:
             return ext_s.Package.model_validate_json(payload).deserialize()
         case EnvelopeFormat.MODEL | EnvelopeFormat.MODEL_WITH_EXTS:
-            msg = "Decoding HUGR envelopes in MODULE format is not supported yet."
-            raise ValueError(msg)
+            # TODO Going via JSON is a temporary solution, until we get model import to
+            # python properly implemented.
+            # https://github.com/CQCL/hugr/issues/2287
+            json_data = rust.to_json_envelope(envelope)
+            return read_envelope(json_data)
 
 
 def read_envelope_hugr(envelope: bytes) -> Hugr:

hugr-py/tests/test_envelope.py

@@ -4,11 +4,14 @@
 import semver
 
 from hugr import ops, tys
+from hugr.build.dfg import Dfg
 from hugr.build.function import Module
 from hugr.envelope import EnvelopeConfig, EnvelopeFormat
 from hugr.hugr.node_port import Node
 from hugr.package import Package
 
+from .conftest import QUANTUM_EXT, H
+
 
 @pytest.fixture
 def package() -> Package:
@@ -29,7 +32,11 @@ def package() -> Package:
 
 def test_envelope(package: Package):
     # Binary compression roundtrip
-    for format in [EnvelopeFormat.JSON]:
+    for format in [
+        EnvelopeFormat.JSON,
+        EnvelopeFormat.MODEL,
+        EnvelopeFormat.MODEL_WITH_EXTS,
+    ]:
         for compression in [None, 0]:
             encoded = package.to_bytes(EnvelopeConfig(format=format, zstd=compression))
             decoded = Package.from_bytes(encoded)
@@ -63,3 +70,15 @@ def test_legacy_funcdefn():
     op2 = h[Node(2)].op
     assert isinstance(op2, ops.FuncDefn)
     assert op2.visibility == "Private"
+
+
+def test_model_import_with_ext():
+    dfg = Dfg(tys.Qubit)
+    h_outs = dfg.add_op(H, dfg.inputs()[0])
+    dfg.set_outputs(h_outs)
+    pkg = Package(modules=[dfg.hugr], extensions=[QUANTUM_EXT])
+    data = pkg.to_bytes(config=EnvelopeConfig.BINARY)
+    pkg1 = Package.from_bytes(data)
+    data1 = pkg1.to_bytes(config=EnvelopeConfig.BINARY)
+    pkg2 = Package.from_bytes(data1)
+    assert pkg2.modules[0].num_nodes() == 8