feat: Add TypeArg/Type::used_extensions

Author: aborgna-q

hugr-py/src/hugr/ext.py

@@ -415,6 +415,14 @@ class ExtensionExists(Exception):
 
         extension_id: ExtensionId
 
+    def ids(self) -> set[ExtensionId]:
+        """Get the set of extension IDs in the registry.
+
+        Returns:
+            Set of extension IDs.
+        """
+        return set(self.extensions.keys())
+
     def add_extension(self, extension: Extension) -> Extension:
         """Add an extension to the registry.
 
@@ -448,3 +456,20 @@ def get_extension(self, name: ExtensionId) -> Extension:
             return self.extensions[name]
         except KeyError as e:
             raise self.ExtensionNotFound(name) from e
+
+    def extend(self, other: ExtensionRegistry) -> None:
+        """Add a registry of extensions to this registry.
+
+        If an extension with the same name already exists, the one with the
+        higher version is kept.
+
+        Args:
+            other: The extension registry to add.
+        """
+        for name, ext in other.extensions.items():
+            if name in self.extensions and self.extensions[name].version >= ext.version:
+                continue
+            self.extensions[name] = ext
+
+    def __str__(self) -> str:
+        return "ExtensionRegistry(" + ", ".join(self.extensions.keys()) + ")"

hugr-py/src/hugr/tys.py

@@ -13,6 +13,7 @@
     from collections.abc import Iterable, Sequence
 
     from hugr import ext
+    from hugr.ext import ExtensionRegistry
 
 
 ExtensionId = stys.ExtensionId
@@ -55,6 +56,21 @@ def to_model(self) -> model.Term | model.Splice:
         """Convert the type argument to a model Term."""
         raise NotImplementedError(self)
 
+    def used_extensions(self) -> ExtensionRegistry:
+        """Get the set of extensions required to define this type argument.
+
+        Raises:
+            UnknownTypeExtensionError: if a type argument contains is a
+                :class:`Opaque` type that has not been resolved.
+
+        Example:
+            >>> TypeTypeArg(ty=Qubit).used_extensions().ids()
+            {'prelude'}
+        """
+        from hugr.ext import ExtensionRegistry
+
+        return ExtensionRegistry()
+
 
 @runtime_checkable
 class Type(Protocol):
@@ -95,10 +111,34 @@ def to_model(self) -> model.Term | model.Splice:
         """Convert the type to a model Term."""
         raise NotImplementedError(self)
 
+    def used_extensions(self) -> ExtensionRegistry:
+        """Get the set of extensions required to define this type.
+
+        Note that :class:`Opaque` types do not know their extension, so they
+        will raise an error. Use :meth:`resolve` to get the actual type
+        and then call this method.
+
+        Raises:
+            UnknownTypeExtensionError: if the type is an :class:`Opaque` type
+                and has not been resolved.
+
+        Example:
+            >>> Qubit.used_extensions().ids()
+            {'prelude'}
+        """
+        from hugr.ext import ExtensionRegistry
+
+        return ExtensionRegistry()
+
 
 #: Row of types.
 TypeRow = list[Type]
 
+
+class UnknownTypeExtensionError(Exception):
+    """Exception raised when querying the extension of an :method:`Opaque` type."""
+
+
 # --------------------------------------------
 # --------------- TypeParam ------------------
 # --------------------------------------------
@@ -211,6 +251,9 @@ def __str__(self) -> str:
     def to_model(self) -> model.Term | model.Splice:
         return self.ty.to_model()
 
+    def used_extensions(self) -> ExtensionRegistry:
+        return self.ty.used_extensions()
+
 
 @dataclass(frozen=True)
 class BoundedNatArg(TypeArg):
@@ -264,6 +307,12 @@ def to_model(self) -> model.Term:
         # For now we assume that this is a list.
         return model.List([elem.to_model() for elem in self.elems])
 
+    def used_extensions(self) -> ExtensionRegistry:
+        reg = super().used_extensions()
+        for arg in self.elems:
+            reg.extend(arg.used_extensions())
+        return reg
+
 
 @dataclass(frozen=True)
 class VariableArg(TypeArg):
@@ -324,6 +373,13 @@ def to_model(self) -> model.Term:
         )
         return model.Apply("core.adt", [variants])
 
+    def used_extensions(self) -> ExtensionRegistry:
+        types = [ty for row in self.variant_rows for ty in row]
+        reg = super().used_extensions()
+        for ty in types:
+            reg.extend(ty.used_extensions())
+        return reg
+
 
 @dataclass(eq=False)
 class UnitSum(Sum):
@@ -457,6 +513,13 @@ def __repr__(self) -> str:
     def to_model(self) -> model.Term:
         return model.Apply("prelude.usize")
 
+    def used_extensions(self) -> ExtensionRegistry:
+        from hugr.std.prelude import PRELUDE_EXTENSION
+
+        reg = super().used_extensions()
+        reg.add_extension(PRELUDE_EXTENSION)
+        return reg
+
 
 @dataclass(frozen=True)
 class Alias(Type):
@@ -543,6 +606,14 @@ def to_model(self) -> model.Term:
         outputs = model.List([output.to_model() for output in self.output])
         return model.Apply("core.fn", [inputs, outputs])
 
+    def used_extensions(self) -> ExtensionRegistry:
+        reg = super().used_extensions()
+        for ty in self.input:
+            reg.extend(ty.used_extensions())
+        for ty in self.output:
+            reg.extend(ty.used_extensions())
+        return reg
+
 
 @dataclass(frozen=True)
 class PolyFuncType(Type):
@@ -587,6 +658,9 @@ def to_model(self) -> model.Term:
         error = "PolyFuncType used as a Type"
         raise TypeError(error)
 
+    def used_extensions(self) -> ExtensionRegistry:
+        return self.body.used_extensions()
+
 
 @dataclass
 class ExtType(Type):
@@ -632,7 +706,7 @@ def __eq__(self, value):
         return super().__eq__(value)
 
     def to_model(self) -> model.Term:
-        # This cast is only neccessary because `Type` can both be an
+        # This cast is only necessary because `Type` can both be an
         # actual type or a row variable.
         args = [cast(model.Term, arg.to_model()) for arg in self.args]
 
@@ -642,6 +716,11 @@ def to_model(self) -> model.Term:
 
         return model.Apply(name, args)
 
+    def used_extensions(self) -> ExtensionRegistry:
+        reg = super().used_extensions()
+        reg.add_extension(self.type_def.get_extension())
+        return reg
+
 
 def _type_str(name: str, args: Sequence[TypeArg]) -> str:
     if len(args) == 0:
@@ -693,6 +772,10 @@ def to_model(self) -> model.Term:
 
         return model.Apply(self.id, args)
 
+    def used_extensions(self) -> ExtensionRegistry:
+        msg = "Opaque types do not know their extension. Call `resolve` first."
+        raise UnknownTypeExtensionError(msg)
+
 
 @dataclass
 class _QubitDef(Type):
@@ -708,6 +791,13 @@ def __repr__(self) -> str:
     def to_model(self) -> model.Term:
         return model.Apply("prelude.qubit", [])
 
+    def used_extensions(self) -> ExtensionRegistry:
+        from hugr.std.prelude import PRELUDE_EXTENSION
+
+        reg = super().used_extensions()
+        reg.add_extension(PRELUDE_EXTENSION)
+        return reg
+
 
 #: Qubit type.
 Qubit = _QubitDef()