fix(hugr-py): ops require their own extensions

Fixes #1301
CQCL · Jul 15, 2024 · c3339fb · c3339fb
1 parent cfb0674
commit c3339fb
Show file tree

Hide file tree

Showing 4 changed files with 23 additions and 8 deletions.
diff --git a/hugr-py/src/hugr/std/int.py b/hugr-py/src/hugr/std/int.py
@@ -65,7 +65,9 @@ def to_custom(self) -> Custom:
         return Custom(
             "idivmod_u",
             tys.FunctionType(
-                input=[int_t(self.arg1)] * 2, output=[int_t(self.arg2)] * 2
+                input=[int_t(self.arg1)] * 2,
+                output=[int_t(self.arg2)] * 2,
+                extension_reqs=[OPS_EXTENSION],
             ),
             extension=OPS_EXTENSION,
             args=[tys.BoundedNatArg(n=self.arg1), tys.BoundedNatArg(n=self.arg2)],

diff --git a/hugr-py/src/hugr/std/logic.py b/hugr-py/src/hugr/std/logic.py
@@ -20,7 +20,8 @@ class _NotDef(AsCustomOp):
     """Not operation."""
 
     def to_custom(self) -> Custom:
-        return Custom("Not", tys.FunctionType.endo([tys.Bool]), extension=EXTENSION_ID)
+        sig = tys.FunctionType.endo([tys.Bool], [EXTENSION_ID])
+        return Custom("Not", sig, extension=EXTENSION_ID)
 
     def __call__(self, a: ComWire) -> Command:
         return DataflowOp.__call__(self, a)

diff --git a/hugr-py/src/hugr/tys.py b/hugr-py/src/hugr/tys.py
@@ -330,14 +330,18 @@ def empty(cls) -> FunctionType:
         return cls(input=[], output=[])
 
     @classmethod
-    def endo(cls, tys: TypeRow) -> FunctionType:
+    def endo(
+        cls, tys: TypeRow, extension_reqs: ExtensionSet | None = None
+    ) -> FunctionType:
         """Function type with the same input and output types.
 
         Example:
             >>> FunctionType.endo([Qubit])
             FunctionType([Qubit], [Qubit])
         """
-        return cls(input=tys, output=tys)
+        return cls(
+            input=tys, output=tys, extension_reqs=extension_reqs or ExtensionSet()
+        )
 
     def flip(self) -> FunctionType:
         """Return a new function type with input and output types swapped.

diff --git a/hugr-py/tests/conftest.py b/hugr-py/tests/conftest.py
@@ -48,7 +48,7 @@ def __call__(self, q: ComWire) -> Command:
     def to_custom(self) -> Custom:
         return Custom(
             self._enum.value,
-            tys.FunctionType.endo([tys.Qubit]),
+            tys.FunctionType.endo([tys.Qubit], extension_reqs=[QUANTUM_EXTENSION_ID]),
             extension=QUANTUM_EXTENSION_ID,
         )
 
@@ -70,7 +70,9 @@ class _Enum(Enum):
     def to_custom(self) -> Custom:
         return Custom(
             self._enum.value,
-            tys.FunctionType.endo([tys.Qubit] * 2),
+            tys.FunctionType.endo(
+                [tys.Qubit] * 2, extension_reqs=[QUANTUM_EXTENSION_ID]
+            ),
             extension=QUANTUM_EXTENSION_ID,
         )
 
@@ -90,7 +92,11 @@ class MeasureDef(AsCustomOp):
     def to_custom(self) -> Custom:
         return Custom(
             "Measure",
-            tys.FunctionType([tys.Qubit], [tys.Qubit, tys.Bool]),
+            tys.FunctionType(
+                [tys.Qubit],
+                [tys.Qubit, tys.Bool],
+                extension_reqs=[QUANTUM_EXTENSION_ID],
+            ),
             extension=QUANTUM_EXTENSION_ID,
         )
 
@@ -106,7 +112,9 @@ class RzDef(AsCustomOp):
     def to_custom(self) -> Custom:
         return Custom(
             "Rz",
-            tys.FunctionType([tys.Qubit, FLOAT_T], [tys.Qubit]),
+            tys.FunctionType(
+                [tys.Qubit, FLOAT_T], [tys.Qubit], extension_reqs=[QUANTUM_EXTENSION_ID]
+            ),
             extension=QUANTUM_EXTENSION_ID,
         )