Support for writing hugr-model envelopes in python.

Author: zrho

hugr-py/src/hugr/envelope.py

@@ -35,6 +35,7 @@
 from dataclasses import dataclass
 from enum import Enum
 from typing import TYPE_CHECKING, ClassVar
+import json
 
 import zstd
 
@@ -53,13 +54,19 @@ def make_envelope(package: Package, config: EnvelopeConfig) -> bytes:
     payload: bytes
     match config.format:
         case EnvelopeFormat.JSON:
-            json = package._to_serial().model_dump_json()
+            json_str = package._to_serial().model_dump_json()
             # This introduces an extra encode/decode roundtrip when calling
             # `make_envelope_str`, but we prioritize speed for binary formats.
-            payload = json.encode("utf-8")
-        case EnvelopeFormat.MODULE | EnvelopeFormat.MODULE_WITH_EXTS:
-            msg = "MODULE encoding of HUGR envelopes is not supported yet."
-            raise ValueError(msg)
+            payload = json_str.encode("utf-8")
+
+        case EnvelopeFormat.MODULE:
+            return bytes(package.to_model())
+
+        case EnvelopeFormat.MODULE_WITH_EXTS:
+            package_bytes = bytes(package.to_model())
+            extension_str = json.dumps([ext.model_dump() for ext in package.extensions])
+            extension_bytes = extension_str.encode("utf8")
+            return package_bytes + extension_bytes
 
     if config.zstd is not None:
         payload = zstd.compress(payload, config.zstd)

hugr-py/src/hugr/package.py

@@ -25,6 +25,7 @@ def _deprecated(func):
     read_envelope_str,
 )
 from hugr.ops import FuncDecl, FuncDefn, Op
+import hugr.model as model
 
 if TYPE_CHECKING:
     from hugr.ext import Extension
@@ -127,6 +128,13 @@ def from_json(cls, json_str: str) -> Package:
         """
         return ext_s.Package.model_validate_json(json_str).deserialize()
 
+    def to_model(self) -> model.Package:
+        """Export the package as its hugr model representation.
+
+        At the moment this does not yet contain the extensions.
+        """
+        return model.Package([module.to_model() for module in self.modules])
+
 
 @dataclass(frozen=True)
 class PackagePointer: