feat: Package pydantic model for modules + extensions (#1387)

Closes #1358
CQCL · Aug 1, 2024 · 68cfac5 · 68cfac5
1 parent dd1dc48
commit 68cfac5
Show file tree

Hide file tree

Showing 11 changed files with 388 additions and 42 deletions.
diff --git a/hugr-py/src/hugr/__init__.py b/hugr-py/src/hugr/__init__.py
@@ -7,6 +7,6 @@
 __version__ = "0.5.0"
 
 
-def get_serialisation_version() -> str:
+def get_serialization_version() -> str:
     """Return the current version of the serialization schema."""
     return "live"
diff --git a/hugr-py/src/hugr/serialization/extension.py b/hugr-py/src/hugr/serialization/extension.py
@@ -3,9 +3,10 @@
 import pydantic as pd
 from pydantic_extra_types.semantic_version import SemanticVersion
 
-from hugr import get_serialisation_version
+from hugr import get_serialization_version
 
 from .ops import Value
+from .serial_hugr import SerialHugr
 from .tys import (
     ConfiguredBaseModel,
     ExtensionId,
@@ -75,4 +76,13 @@ class Extension(ConfiguredBaseModel):
 
     @classmethod
     def get_version(cls) -> str:
-        return get_serialisation_version()
+        return get_serialization_version()
+
+
+class Package(ConfiguredBaseModel):
+    modules: list[SerialHugr]
+    extensions: list[Extension] = pd.Field(default_factory=list)
+
+    @classmethod
+    def get_version(cls) -> str:
+        return get_serialization_version()
diff --git a/hugr-py/src/hugr/serialization/serial_hugr.py b/hugr-py/src/hugr/serialization/serial_hugr.py
@@ -3,7 +3,7 @@
 from pydantic import ConfigDict, Field
 
 import hugr
-from hugr import get_serialisation_version
+from hugr import get_serialization_version
 from hugr.node_port import NodeIdx, PortOffset
 
 from .ops import OpType
@@ -14,7 +14,7 @@
 Edge = tuple[Port, Port]
 
 VersionField = Field(
-    default_factory=get_serialisation_version,
+    default_factory=get_serialization_version,
     title="Version",
     description="Serialisation Schema Version",
     frozen=True,

diff --git a/hugr-py/tests/serialization/test_basic.py b/hugr-py/tests/serialization/test_basic.py
@@ -1,11 +1,11 @@
-from hugr import get_serialisation_version
+from hugr import get_serialization_version
 from hugr.serialization.serial_hugr import SerialHugr
 
 
 def test_empty():
     h = SerialHugr(nodes=[], edges=[])
     assert h.model_dump() == {
-        "version": get_serialisation_version(),
+        "version": get_serialization_version(),
         "nodes": [],
         "edges": [],
         "metadata": None,

diff --git a/hugr-py/tests/serialization/test_extension.py b/hugr-py/tests/serialization/test_extension.py
@@ -1,12 +1,15 @@
 from semver import Version
 
+from hugr import get_serialization_version
 from hugr.serialization.extension import (
     ExplicitBound,
     Extension,
     OpDef,
+    Package,
     TypeDef,
     TypeDefBound,
 )
+from hugr.serialization.serial_hugr import SerialHugr
 from hugr.serialization.tys import (
     FunctionType,
     PolyFuncType,
@@ -71,7 +74,8 @@
 """
 
 
-def test_deserialize():
+def test_extension():
+    assert get_serialization_version() == Extension.get_version()
     param = TypeParam(root=TypeTypeParam(b=TypeBound.Copyable))
 
     bound = TypeDefBound(root=ExplicitBound(bound=TypeBound.Copyable))
@@ -106,3 +110,23 @@ def test_deserialize():
     dumped_json = ext.model_dump_json()
 
     assert Extension.model_validate_json(dumped_json) == ext
+
+
+def test_package():
+    assert get_serialization_version() == Package.get_version()
+
+    ext = Extension(
+        version=Version(0, 1, 0),
+        name="ext",
+        extension_reqs=set(),
+        types={},
+        values={},
+        operations={},
+    )
+    ext_load = Extension.model_validate_json(EXAMPLE)
+    package = Package(
+        extensions=[ext, ext_load], modules=[SerialHugr(nodes=[], edges=[])]
+    )
+
+    package_load = Package.model_validate_json(package.model_dump_json())
+    assert package == package_load
diff --git a/poetry.lock b/poetry.lock
diff --git a/scripts/generate_schema.py b/scripts/generate_schema.py
@@ -16,7 +16,7 @@
 from pydantic import ConfigDict
 from pydantic.json_schema import models_json_schema
 
-from hugr.serialization.extension import Extension
+from hugr.serialization.extension import Extension, Package
 from hugr.serialization.serial_hugr import SerialHugr
 from hugr.serialization.testing_hugr import TestingHugr
 
@@ -33,7 +33,7 @@ def write_schema(
     path = out_dir / filename
     print(f"Rebuilding model with config: {config}")
     schema._pydantic_rebuild(config or ConfigDict(), force=True, **kwargs)
-    schemas = [schema, Extension]
+    schemas = [schema, Extension, Package]
     print(f"Writing schema to {path}")
     _, top_level_schema = models_json_schema(
         [(s, "validation") for s in schemas], title="HUGR schema"

diff --git a/specification/schema/hugr_schema_live.json b/specification/schema/hugr_schema_live.json
@@ -1372,6 +1372,29 @@
             "title": "Output",
             "type": "object"
         },
+        "Package": {
+            "properties": {
+                "modules": {
+                    "items": {
+                        "$ref": "#/$defs/SerialHugr"
+                    },
+                    "title": "Modules",
+                    "type": "array"
+                },
+                "extensions": {
+                    "items": {
+                        "$ref": "#/$defs/Extension"
+                    },
+                    "title": "Extensions",
+                    "type": "array"
+                }
+            },
+            "required": [
+                "modules"
+            ],
+            "title": "Package",
+            "type": "object"
+        },
         "PolyFuncType": {
             "additionalProperties": true,
             "description": "A polymorphic type scheme, i.e. of a FuncDecl, FuncDefn or OpDef.  (Nodes/operations in the Hugr are not polymorphic.)",

diff --git a/specification/schema/hugr_schema_strict_live.json b/specification/schema/hugr_schema_strict_live.json
@@ -1372,6 +1372,29 @@
             "title": "Output",
             "type": "object"
         },
+        "Package": {
+            "properties": {
+                "modules": {
+                    "items": {
+                        "$ref": "#/$defs/SerialHugr"
+                    },
+                    "title": "Modules",
+                    "type": "array"
+                },
+                "extensions": {
+                    "items": {
+                        "$ref": "#/$defs/Extension"
+                    },
+                    "title": "Extensions",
+                    "type": "array"
+                }
+            },
+            "required": [
+                "modules"
+            ],
+            "title": "Package",
+            "type": "object"
+        },
         "PolyFuncType": {
             "additionalProperties": false,
             "description": "A polymorphic type scheme, i.e. of a FuncDecl, FuncDefn or OpDef.  (Nodes/operations in the Hugr are not polymorphic.)",