feat(hugr-py): Add to/from_bytes/str to Hugr, using envelopes (#2228)

Followup to #2148

- Adds envelope reading/writting methods to `Hugr` (py) that mimic
Package.
- Deprecates `to_json` and `load_json`.
- Uses envelopes in python tests.

Note that `EnvelopeConfig.BINARY` is still set to JSON.
Using `hugr-module` there currently causes errors.

Author: aborgna-q

hugr-py/src/hugr/envelope.py

@@ -40,16 +40,22 @@
 import pyzstd
 
 if TYPE_CHECKING:
+    from hugr.hugr.base import Hugr
     from hugr.package import Package
 
 # This is a hard-coded magic number that identifies the start of a HUGR envelope.
 MAGIC_NUMBERS = b"HUGRiHJv"
 
 
-def make_envelope(package: Package, config: EnvelopeConfig) -> bytes:
-    """Encode a HUGR package into an envelope, using the given configuration."""
+def make_envelope(package: Package | Hugr, config: EnvelopeConfig) -> bytes:
+    """Encode a HUGR or Package into an envelope, using the given configuration."""
+    from hugr.package import Package
+
     envelope = bytearray(config._make_header().to_bytes())
 
+    if not isinstance(package, Package):
+        package = Package(modules=[package], extensions=[])
+
     # Currently only uncompressed JSON is supported.
     payload: bytes
     match config.format:
@@ -77,8 +83,8 @@ def make_envelope(package: Package, config: EnvelopeConfig) -> bytes:
     return bytes(envelope)
 
 
-def make_envelope_str(package: Package, config: EnvelopeConfig) -> str:
-    """Encode a HUGR package into an envelope, using the given configuration."""
+def make_envelope_str(package: Package | Hugr, config: EnvelopeConfig) -> str:
+    """Encode a HUGR or Package into an envelope, using the given configuration."""
     if not config.format.ascii_printable():
         msg = "Only ascii-printable envelope formats can be encoded into a string."
         raise ValueError(msg)
@@ -104,11 +110,43 @@ def read_envelope(envelope: bytes) -> Package:
             raise ValueError(msg)
 
 
+def read_envelope_hugr(envelope: bytes) -> Hugr:
+    """Decode a HUGR from an envelope.
+
+    Raises:
+        ValueError: If the envelope does not contain a single module.
+    """
+    pkg = read_envelope(envelope)
+    if len(pkg.modules) != 1:
+        msg = (
+            "Expected a single module in the envelope, but got "
+            + f"{len(pkg.modules)} modules."
+        )
+        raise ValueError(msg)
+    return pkg.modules[0]
+
+
 def read_envelope_str(envelope: str) -> Package:
     """Decode a HUGR package from an envelope."""
     return read_envelope(envelope.encode("utf-8"))
 
 
+def read_envelope_hugr_str(envelope: str) -> Hugr:
+    """Decode a HUGR from an envelope.
+
+    Raises:
+        ValueError: If the envelope does not contain a single module.
+    """
+    pkg = read_envelope_str(envelope)
+    if len(pkg.modules) != 1:
+        msg = (
+            "Expected a single module in the envelope, but got "
+            + f"{len(pkg.modules)} modules."
+        )
+        raise ValueError(msg)
+    return pkg.modules[0]
+
+
 class EnvelopeFormat(Enum):
     """Format used to encode a HUGR envelope."""
 

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

@@ -15,10 +15,19 @@
     overload,
 )
 
+from typing_extensions import deprecated
+
 import hugr.model as model
 import hugr.ops as ops
 from hugr._serialization.ops import OpType as SerialOp
 from hugr._serialization.serial_hugr import SerialHugr
+from hugr.envelope import (
+    EnvelopeConfig,
+    make_envelope,
+    make_envelope_str,
+    read_envelope_hugr,
+    read_envelope_hugr_str,
+)
 from hugr.exceptions import ParentBeforeChild
 from hugr.ops import (
     CFG,
@@ -892,6 +901,59 @@ def get_meta(idx: int) -> dict[str, Any]:
 
         return hugr
 
+    @staticmethod
+    def from_bytes(envelope: bytes) -> Hugr:
+        """Deserialize a byte string to a Hugr object.
+
+        Some envelope formats can be read from a string. See :meth:`from_str`.
+
+        Args:
+            envelope: The byte string representing a Hugr envelope.
+
+        Returns:
+            The deserialized Hugr object.
+
+        Raises:
+            ValueError: If the envelope does not contain exactly one module.
+        """
+        return read_envelope_hugr(envelope)
+
+    @staticmethod
+    def from_str(envelope: str) -> Hugr:
+        """Deserialize a string to a Hugr object.
+
+        Not all envelope formats can be read from a string.
+        See :meth:`from_bytes` for a more general method.
+
+        Args:
+            envelope: The string representing a Hugr envelope.
+
+        Returns:
+            The deserialized Hugr object.
+
+        Raises:
+            ValueError: If the envelope does not contain exactly one module.
+        """
+        return read_envelope_hugr_str(envelope)
+
+    def to_bytes(self, config: EnvelopeConfig | None = None) -> bytes:
+        """Serialize the HUGR into an envelope byte string.
+
+        Some envelope formats can be encoded into a string. See :meth:`to_str`.
+        """
+        config = config or EnvelopeConfig.BINARY
+        return make_envelope(self, config)
+
+    def to_str(self, config: EnvelopeConfig | None = None) -> str:
+        """Serialize the package to a HUGR envelope string.
+
+        Not all envelope formats can be encoded into a string.
+        See :meth:`to_bytes` for a more general method.
+        """
+        config = config or EnvelopeConfig.TEXT
+        return make_envelope_str(self, config)
+
+    @deprecated("Use HUGR envelopes instead. See the `to_bytes` and `to_str` methods.")
     def to_json(self) -> str:
         """Serialize the HUGR to a JSON string.
 
@@ -909,6 +971,7 @@ def to_model(self) -> model.Module:
         return model.Module(region)
 
     @classmethod
+    @deprecated("Use HUGR envelopes instead. See the `to_bytes` and `to_str` methods.")
     def load_json(cls, json_str: str) -> Hugr:
         """Deserialize a JSON string into a HUGR.
 

hugr-py/tests/conftest.py

@@ -1,6 +1,5 @@
 from __future__ import annotations
 
-import json
 import os
 import pathlib
 import subprocess
@@ -11,7 +10,6 @@
 from typing_extensions import Self
 
 from hugr import ext, tys
-from hugr._serialization.serial_hugr import SerialHugr
 from hugr.envelope import EnvelopeConfig
 from hugr.hugr import Hugr
 from hugr.ops import AsExtOp, Command, DataflowOp, ExtOp, RegisteredOp
@@ -131,7 +129,7 @@ def _base_command() -> list[str]:
 def mermaid(h: Hugr):
     """Render the Hugr as a mermaid diagram for debugging."""
     cmd = [*_base_command(), "mermaid", "-"]
-    _run_hugr_cmd(h.to_json().encode(), cmd)
+    _run_hugr_cmd(h.to_str().encode(), cmd)
 
 
 def validate(
@@ -151,21 +149,17 @@ def validate(
     # TODO: Use envelopes instead of legacy hugr-json
     cmd = [*_base_command(), "validate", "-"]
 
-    if isinstance(h, Hugr):
-        cmd += ["--hugr-json"]
-        _run_hugr_cmd(h.to_json().encode(), cmd)
-    else:
-        serial = h.to_bytes(EnvelopeConfig.BINARY)
-        _run_hugr_cmd(serial, cmd)
+    serial = h.to_bytes(EnvelopeConfig.BINARY)
+    _run_hugr_cmd(serial, cmd)
 
     if not roundtrip:
         return
 
     # Roundtrip checks
     if isinstance(h, Hugr):
-        starting_json = json.loads(h.to_json())
-        h2 = Hugr._from_serial(SerialHugr.load_json(starting_json))
-        roundtrip_json = json.loads(h2._to_serial().to_json())
+        starting_json = h.to_str()
+        h2 = Hugr.from_str(starting_json)
+        roundtrip_json = h2.to_str()
         assert roundtrip_json == starting_json
 
         if snap is not None:

hugr-py/tests/serialization/test_basic.py

@@ -23,7 +23,7 @@ def test_children():
     h = mod.hugr
     assert len(h.children()) == 1
 
-    h2 = Hugr.load_json(h.to_json())
+    h2 = Hugr.from_str(h.to_str())
 
     assert len(h2.children()) == 1
 
@@ -44,7 +44,7 @@ def test_entrypoint():
     assert h[func].parent == h.module_root
 
     # Do a roundtrip, and test all again
-    h2 = Hugr.load_json(h.to_json())
+    h2 = Hugr.from_str(h.to_str())
 
     dfg = h2.entrypoint
     assert h2[dfg].op == ops.DFG(inputs=[tys.Bool], _outputs=[tys.Bool])