From e7f9f4c739e1b523fb7f2b7df0246baea2abf3df Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Agust=C3=ADn=20Borgna?=
 <121866228+aborgna-q@users.noreply.github.com>
Date: Mon, 12 Aug 2024 10:47:58 +0100
Subject: [PATCH] feat(py): `Hugr.to_json` and `.load_json` helpers (#1403)

drive-by: For some reason mypy started complaining about some unrelated
imports, this includes a fix -.-'
---
 .pre-commit-config.yaml          |  2 +-
 hugr-py/src/hugr/hugr.py         | 12 ++++++++++++
 hugr-py/tests/conftest.py        |  2 +-
 hugr-py/tests/test_hugr_build.py | 14 +++++++++++++-
 4 files changed, 27 insertions(+), 3 deletions(-)

diff --git a/.pre-commit-config.yaml b/.pre-commit-config.yaml
index 78b9baff7..a5b43f7ab 100644
--- a/.pre-commit-config.yaml
+++ b/.pre-commit-config.yaml
@@ -40,7 +40,7 @@ repos:
       - id: ruff-format
 
   - repo: https://github.com/pre-commit/mirrors-mypy
-    rev: v1.9.0
+    rev: v1.11.0
     hooks:
       - id: mypy
         additional_dependencies: [pydantic]
diff --git a/hugr-py/src/hugr/hugr.py b/hugr-py/src/hugr/hugr.py
index 99392e15e..5918fed5b 100644
--- a/hugr-py/src/hugr/hugr.py
+++ b/hugr-py/src/hugr/hugr.py
@@ -2,6 +2,7 @@
 
 from __future__ import annotations
 
+import json
 from collections.abc import Iterable, Iterator, Mapping
 from dataclasses import dataclass, field, replace
 from typing import (
@@ -623,3 +624,14 @@ def from_serial(cls, serial: SerialHugr) -> Hugr:
             )
 
         return hugr
+
+    def to_json(self) -> str:
+        """Serialize the HUGR to a JSON string."""
+        return self.to_serial().to_json()
+
+    @classmethod
+    def load_json(cls, json_str: str) -> Hugr:
+        """Deserialize a JSON string into a HUGR."""
+        json_dict = json.loads(json_str)
+        serial = SerialHugr.load_json(json_dict)
+        return cls.from_serial(serial)
diff --git a/hugr-py/tests/conftest.py b/hugr-py/tests/conftest.py
index 94698f9b9..d54302a42 100644
--- a/hugr-py/tests/conftest.py
+++ b/hugr-py/tests/conftest.py
@@ -10,7 +10,7 @@
 
 from typing_extensions import Self
 
-from hugr import tys
+import hugr.tys as tys
 from hugr.hugr import Hugr
 from hugr.ops import AsCustomOp, Command, Custom, DataflowOp
 from hugr.serialization.serial_hugr import SerialHugr
diff --git a/hugr-py/tests/test_hugr_build.py b/hugr-py/tests/test_hugr_build.py
index 021cfc5cf..5ae308f26 100644
--- a/hugr-py/tests/test_hugr_build.py
+++ b/hugr-py/tests/test_hugr_build.py
@@ -2,7 +2,9 @@
 
 import pytest
 
-from hugr import ops, tys, val
+import hugr.ops as ops
+import hugr.tys as tys
+import hugr.val as val
 from hugr.dfg import Dfg, _ancestral_sibling
 from hugr.function import Module
 from hugr.hugr import Hugr
@@ -67,6 +69,16 @@ def test_simple_id():
     validate(simple_id().hugr)
 
 
+def test_json_roundtrip():
+    hugr = simple_id().hugr
+    json = hugr.to_json()
+
+    hugr2 = Hugr.load_json(json)
+    json2 = hugr2.to_json()
+
+    assert json2 == json
+
+
 def test_multiport():
     h = Dfg(tys.Bool)
     (a,) = h.inputs()