feat: Add debug info metadata specification in hugr-py

hugr-py/rust/metadata.rs

@@ -22,4 +22,8 @@ pub mod metadata {
 
     #[pymodule_export]
     const HUGR_USED_EXTENSIONS: &str = hugr_core::metadata::HugrUsedExtensions::KEY;
+
+    // TODO: Get from rust implementation once it exists.
+    #[pymodule_export]
+    const HUGR_DEBUG_INFO: &str = "core.debug_info";
 }

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

@@ -1,2 +1,3 @@
 HUGR_GENERATOR: str
 HUGR_USED_EXTENSIONS: str
+HUGR_DEBUG_INFO: str

hugr-py/src/hugr/debug_info.py

@@ -0,0 +1,152 @@
+"""Typed generator source debug information metadata for HUGR nodes."""
+
+from abc import ABC, abstractmethod
+from dataclasses import dataclass
+from typing import ClassVar, cast
+
+from hugr.utils import JsonType
+
+
+@dataclass
+class DebugRecord(ABC):
+    """Abstract base class for debug records."""
+
+    @abstractmethod
+    def to_json(self) -> JsonType:
+        """Encodes the record as a dictionary of native types that can be serialized by
+        `json.dump`.
+        """
+
+    @classmethod
+    def from_json(cls, value: JsonType) -> "DebugRecord":
+        """Decode a debug record from json. This is not an abstract method because when
+        decoding from json by calling `DebugRecord.from_json` we do not have concrete
+        subtype information, so we decode from the explicit variant tag stored in `kind`
+        instead.
+        """
+        if not isinstance(value, dict):
+            msg = f"Expected a dictionary for DebugRecord, but got {type(value)}"
+            raise TypeError(msg)
+
+        kind = value.get("kind")
+        if isinstance(kind, str):
+            if kind == DICompileUnit.KIND:
+                return DICompileUnit.from_json(value)
+            if kind == DISubprogram.KIND:
+                return DISubprogram.from_json(value)
+            if kind == DILocation.KIND:
+                return DILocation.from_json(value)
+            msg = f"Unknown DebugRecord kind: {kind}"
+            raise TypeError(msg)
+
+        msg = "Expected DebugRecord to contain string field 'kind'."
+        raise TypeError(msg)
+
+
+@dataclass
+class DICompileUnit(DebugRecord):
+    """Debug information for a compilation unit, corresponds to a HUGR module node."""
+
+    KIND: ClassVar[str] = "compile_unit"
+
+    directory: str  # Working directory of the compiler that generated the HUGR.
+    filename: int  # File that contains the HUGR entrypoint.
+    file_table: list[str]  # Global table of all files referenced in the module.
+
+    def to_json(self) -> dict[str, JsonType]:
+        return {
+            "kind": self.KIND,
+            "directory": self.directory,
+            "filename": self.filename,
+            "file_table": cast("list[JsonType]", self.file_table),
+        }
+
+    @classmethod
+    def from_json(cls, value: JsonType) -> "DICompileUnit":
+        if not isinstance(value, dict):
+            msg = f"Expected a dictionary for DICompileUnit, but got {type(value)}"
+            raise TypeError(msg)
+        for key in ("kind", "directory", "filename", "file_table"):
+            if key not in value:
+                msg = f"Expected DICompileUnit to have a '{key}' key but got {value}"
+                raise TypeError(msg)
+        files = value["file_table"]
+        if not isinstance(files, list):
+            msg = f"Expected 'file_table' to be a list but got {type(files)}"
+            raise TypeError(msg)
+        return DICompileUnit(
+            directory=str(value["directory"]),
+            filename=int(value["filename"]),
+            file_table=list[str](value["file_table"]),
+        )
+
+
+@dataclass
+class DISubprogram(DebugRecord):
+    """Debug information for a subprogram, corresponds to a function definition or
+    declaration node.
+    """
+
+    KIND: ClassVar[str] = "subprogram"
+
+    file: int  # Index into the string table for filenames.
+    line_no: int  # First line of the function definition.
+    scope_line: int | None = None  # First line of the function body.
+
+    def to_json(self) -> dict[str, str]:
+        data = {
+            "kind": self.KIND,
+            "file": str(self.file),
+            "line_no": str(self.line_no),
+        }
+        # Declarations have no function body so could have no scope_line.
+        if self.scope_line is not None:
+            data["scope_line"] = str(self.scope_line)
+        return data
+
+    @classmethod
+    def from_json(cls, value: JsonType) -> "DISubprogram":
+        if not isinstance(value, dict):
+            msg = f"Expected a dictionary for DISubprogram, but got {type(value)}"
+            raise TypeError(msg)
+        for key in ("kind", "file", "line_no"):
+            if key not in value:
+                msg = f"Expected DISubprogram to have a '{key}' key but got {value}"
+                raise TypeError(msg)
+        # Declarations have no function body so could have no scope_line.
+        scope_line = int(value["scope_line"]) if "scope_line" in value else None
+        return DISubprogram(
+            file=int(value["file"]),
+            line_no=int(value["line_no"]),
+            scope_line=scope_line,
+        )
+
+
+@dataclass
+class DILocation(DebugRecord):
+    """Debug information for a location, corresponds to call or extension operation
+    node.
+    """
+
+    KIND: ClassVar[str] = "location"
+
+    column: int
+    line_no: int
+
+    def to_json(self) -> dict[str, str]:
+        return {
+            "kind": self.KIND,
+            "column": str(self.column),
+            "line_no": str(self.line_no),
+        }
+
+    @classmethod
+    def from_json(cls, value: JsonType) -> "DILocation":
+        if not isinstance(value, dict):
+            msg = f"Expected a dictionary for DILocation, but got {type(value)}"
+            raise TypeError(msg)
+        for key in ("kind", "column", "line_no"):
+            if key not in value:
+                msg = f"Expected DILocation to have a '{key}' key but got {value}"
+                raise TypeError(msg)
+        return DILocation(column=int(value["column"]), line_no=int(value["line_no"]))

hugr-py/src/hugr/metadata.py

@@ -15,16 +15,16 @@
 )
 
 import hugr._hugr.metadata as rust_metadata
+from hugr.debug_info import DebugRecord
 from hugr.envelope import ExtensionDesc, GeneratorDesc
 
 if TYPE_CHECKING:
     from collections.abc import Iterable, Iterator
 
-Meta = TypeVar("Meta")
+    from hugr.utils import JsonType
+
 
-# Type alias for json values.
-# See <https://github.com/python/typing/issues/182#issuecomment-1320974824>
-JsonType = str | int | float | bool | None | Mapping[str, "JsonType"] | list["JsonType"]
+Meta = TypeVar("Meta")
 
 
 class Metadata(Protocol[Meta]):
@@ -193,3 +193,17 @@ def from_json(cls, value: JsonType) -> list[ExtensionDesc]:
             )
             raise TypeError(msg)
         return [ExtensionDesc.from_json(e) for e in value]
+
+
+class HugrDebugInfo(Metadata[DebugRecord]):
+    """Metadata storing debug information obtained from the generator source."""
+
+    KEY = rust_metadata.HUGR_DEBUG_INFO
+
+    @classmethod
+    def to_json(cls, value: DebugRecord) -> JsonType:
+        return value.to_json()
+
+    @classmethod
+    def from_json(cls, value: JsonType) -> DebugRecord:
+        return DebugRecord.from_json(value)

hugr-py/src/hugr/utils.py

@@ -4,6 +4,10 @@
 from dataclasses import dataclass, field
 from typing import Generic, Protocol, TypeVar
 
+# Type alias for json values.
+# See <https://github.com/python/typing/issues/182#issuecomment-1320974824>
+JsonType = str | int | float | bool | None | Mapping[str, "JsonType"] | list["JsonType"]
+
 L = TypeVar("L", bound=Hashable)
 R = TypeVar("R", bound=Hashable)
 

hugr-py/tests/test_metadata.py

@@ -4,9 +4,18 @@
 
 from semver import Version
 
+from hugr import ops, tys
+from hugr.build.function import Module
+from hugr.debug_info import DICompileUnit, DILocation, DISubprogram
 from hugr.envelope import EnvelopeConfig, ExtensionDesc, GeneratorDesc
 from hugr.hugr import Hugr
-from hugr.metadata import HugrGenerator, HugrUsedExtensions, JsonType, Metadata
+from hugr.metadata import (
+    HugrDebugInfo,
+    HugrGenerator,
+    HugrUsedExtensions,
+    Metadata,
+)
+from hugr.utils import JsonType
 
 
 class CustomMetadata(Metadata[list[JsonType]]):
@@ -85,3 +94,45 @@ def test_metadata_default() -> None:
     ) == GeneratorDesc("hugr-py-test", Version.parse("1.2.3"))
     assert node.metadata.get("missing.metadata") is None
     assert node.metadata.get("missing.metadata", [1, 2, 3]) == [1, 2, 3]
+
+
+def test_debug_info_roundtrip() -> None:
+    mod = Module()
+
+    # Add DICompileUnit debug info to the module root
+    compile_unit = DICompileUnit(
+        directory="/user/project/",
+        filename=0,
+        file_table=["guppy1.py", "guppy2.py"],
+    )
+    mod.hugr[mod.hugr.module_root].metadata[HugrDebugInfo] = compile_unit
+
+    # Add a FuncDefn node to test DISubprogram debug info
+    func = mod.define_function("random_func", [tys.Bool], [tys.Bool])
+    [b] = func.inputs()
+    func.set_outputs(b)
+    subprogram = DISubprogram(file=0, line_no=10, scope_line=11)
+    mod.hugr[func.parent_node].metadata[HugrDebugInfo] = subprogram
+
+    # Add a call node test DILocation debug info
+    caller = mod.define_function("caller", [tys.Bool], [tys.Bool])
+    [b] = caller.inputs()
+    call_node = caller.call(func.parent_node, b)
+    caller.set_outputs(call_node)
+    location = DILocation(column=5, line_no=20)
+    mod.hugr[call_node].metadata[HugrDebugInfo] = location
+
+    # Roundtrip serialization
+    data = mod.hugr.to_bytes(EnvelopeConfig.TEXT)
+    loaded = Hugr.from_bytes(data)
+    module_node = loaded[loaded.module_root]
+
+    assert module_node.metadata[HugrDebugInfo] == compile_unit
+
+    [func_n, caller_n] = loaded.children(loaded.module_root)
+    assert loaded[func_n].metadata[HugrDebugInfo] == subprogram
+
+    call_n = next(
+        n for n in loaded.children(caller_n) if isinstance(loaded[n].op, ops.Call)
+    )
+    assert loaded[call_n].metadata[HugrDebugInfo] == location