feat: Explicit exports for tket_exts ops and types

tket-exts/src/tket_exts/__init__.py

@@ -1,73 +1,54 @@
 """HUGR extension definitions for tket circuits."""
 
-import pkgutil
-import functools
-
+from tket_exts.tket.bool import BoolExtension
+from tket_exts.tket.debug import DebugExtension
+from tket_exts.tket.guppy import GuppyExtension
+from tket_exts.tket.rotation import RotationExtension
+from tket_exts.tket.futures import FuturesExtension
+from tket_exts.tket.qsystem import (
+    QSystemExtension,
+    QSystemRandomExtension,
+    QSystemUtilsExtension,
+)
+from tket_exts.tket.quantum import QuantumExtension
+from tket_exts.tket.result import ResultExtension
+from tket_exts.tket.wasm import WasmExtension
+
+from typing_extensions import deprecated
 from hugr.ext import Extension
-
+from tket_exts import tket
 
 # This is updated by our release-please workflow, triggered by this
 # annotation: x-release-please-version
 __version__ = "0.10.1"
 
-
-@functools.cache
+__all__ = [
+    "bool",
+    "debug",
+    "guppy",
+    "rotation",
+    "futures",
+    "qsystem",
+    "qsystem_random",
+    "qsystem_utils",
+    "quantum",
+    "result",
+    "wasm",
+]
+
+bool: BoolExtension = tket.bool.BoolExtension()
+debug: DebugExtension = tket.debug.DebugExtension()
+guppy: GuppyExtension = tket.guppy.GuppyExtension()
+rotation: RotationExtension = tket.rotation.RotationExtension()
+futures: FuturesExtension = tket.futures.FuturesExtension()
+qsystem: QSystemExtension = tket.qsystem.QSystemExtension()
+qsystem_random: QSystemRandomExtension = tket.qsystem.QSystemRandomExtension()
+qsystem_utils: QSystemUtilsExtension = tket.qsystem.QSystemUtilsExtension()
+quantum: QuantumExtension = tket.quantum.QuantumExtension()
+result: ResultExtension = tket.result.ResultExtension()
+wasm: WasmExtension = tket.wasm.WasmExtension()
+
+
+@deprecated("Use tket_exts.bool() instead")
 def opaque_bool() -> Extension:
-    return load_extension("tket.bool")
-
-
-@functools.cache
-def debug() -> Extension:
-    return load_extension("tket.debug")
-
-
-@functools.cache
-def guppy() -> Extension:
-    return load_extension("tket.guppy")
-
-
-@functools.cache
-def rotation() -> Extension:
-    return load_extension("tket.rotation")
-
-
-@functools.cache
-def futures() -> Extension:
-    return load_extension("tket.futures")
-
-
-@functools.cache
-def qsystem() -> Extension:
-    return load_extension("tket.qsystem")
-
-
-@functools.cache
-def qsystem_random() -> Extension:
-    return load_extension("tket.qsystem.random")
-
-
-@functools.cache
-def qsystem_utils() -> Extension:
-    return load_extension("tket.qsystem.utils")
-
-
-@functools.cache
-def quantum() -> Extension:
-    return load_extension("tket.quantum")
-
-
-@functools.cache
-def result() -> Extension:
-    return load_extension("tket.result")
-
-
-@functools.cache
-def wasm() -> Extension:
-    return load_extension("tket.wasm")
-
-
-def load_extension(name: str) -> Extension:
-    replacement = name.replace(".", "/")
-    json_str = pkgutil.get_data(__name__, f"data/{replacement}.json")
-    assert json_str is not None, f"Could not load json for extension {name}"
-    return Extension.from_json(json_str.decode())
+    return bool()

tket-exts/src/tket_exts/tket/__init__.py

@@ -0,0 +1,15 @@
+"""HUGR extension definitions for tket."""
+
+from . import bool, debug, guppy, futures, qsystem, quantum, result, rotation, wasm
+
+__all__ = [
+    "bool",
+    "debug",
+    "guppy",
+    "futures",
+    "qsystem",
+    "quantum",
+    "result",
+    "rotation",
+    "wasm",
+]

tket-exts/src/tket_exts/tket/_util.py

@@ -0,0 +1,29 @@
+"""Utility functions for tket extensions."""
+
+import pkgutil
+from typing import List, Protocol
+
+from hugr.ext import Extension, OpDef, TypeDef
+from semver import Version
+
+
+def load_extension(name: str) -> Extension:
+    import tket_exts
+
+    replacement = name.replace(".", "/")
+    json_str = pkgutil.get_data(tket_exts.__name__, f"data/{replacement}.json")
+    assert json_str is not None, f"Could not load json for extension {name}"
+    return Extension.from_json(json_str.decode())
+
+
+class TketExtension(Protocol):
+    """A protocol for tket extensions."""
+
+    def TYPES(self) -> List[TypeDef]: ...
+    def OPS(self) -> List[OpDef]: ...
+    def __call__(self) -> Extension: ...
+
+    @property
+    def version(self) -> Version:
+        """The version of the extension"""
+        return self().version

tket-exts/src/tket_exts/tket/bool.py

@@ -0,0 +1,78 @@
+"""Opaque boolean operations."""
+
+import functools
+from typing import List
+
+from hugr.ext import Extension, OpDef, TypeDef
+from hugr.ops import ExtOp
+from hugr.tys import ExtType
+from ._util import TketExtension, load_extension
+
+
+class BoolExtension(TketExtension):
+    """Opaque boolean extension.
+
+    This extension operates with an opaque `bool_t` type that
+    is equivalent to Hugr's `[] + []` sum type representation.
+    """
+
+    @functools.cache
+    def __call__(self) -> Extension:
+        """Return the bool extension"""
+        return load_extension("tket.bool")
+
+    def TYPES(self) -> List[TypeDef]:
+        """Return the types defined by this extension"""
+        return [self.bool_t.type_def]
+
+    def OPS(self) -> List[OpDef]:
+        """Return the operations defined by this extension"""
+        return [
+            self.and_op.op_def(),
+            self.eq.op_def(),
+            self.make_opaque.op_def(),
+            self.not_op.op_def(),
+            self.or_op.op_def(),
+            self.read.op_def(),
+            self.xor.op_def(),
+        ]
+
+    @functools.cached_property
+    def bool_t(self) -> ExtType:
+        """An opaque boolean type"""
+        return self().get_type("bool").instantiate([])
+
+    @functools.cached_property
+    def and_op(self) -> ExtOp:
+        """Logical AND between two tket.bools"""
+        return self().get_op("and").instantiate()
+
+    @functools.cached_property
+    def eq(self) -> ExtOp:
+        """Equality between two tket.bools"""
+        return self().get_op("eq").instantiate()
+
+    @functools.cached_property
+    def make_opaque(self) -> ExtOp:
+        """Convert a Hugr `bool_t` (a unit sum) into an tket.bool."""
+        return self().get_op("make_opaque").instantiate()
+
+    @functools.cached_property
+    def not_op(self) -> ExtOp:
+        """Negation of a tket.bool"""
+        return self().get_op("not").instantiate()
+
+    @functools.cached_property
+    def or_op(self) -> ExtOp:
+        """Logical OR between two tket.bools"""
+        return self().get_op("or").instantiate()
+
+    @functools.cached_property
+    def read(self) -> ExtOp:
+        """Convert a tket.bool into a Hugr bool_t (a unit sum)"""
+        return self().get_op("read").instantiate()
+
+    @functools.cached_property
+    def xor(self) -> ExtOp:
+        """Logical XOR between two tket.bools"""
+        return self().get_op("xor").instantiate()

tket-exts/src/tket_exts/tket/debug.py

@@ -0,0 +1,47 @@
+"""Debug extension operations."""
+
+import functools
+from typing import List
+
+from hugr.ext import Extension, OpDef, TypeDef
+from hugr.ops import ExtOp
+from hugr.tys import StringArg, BoundedNatArg
+from ._util import TketExtension, load_extension
+
+
+class DebugExtension(TketExtension):
+    """Extension for debugging operations."""
+
+    @functools.cache
+    def __call__(self) -> Extension:
+        """Return the debug extension"""
+        return load_extension("tket.debug")
+
+    def TYPES(self) -> List[TypeDef]:
+        """Return the types defined by this extension"""
+        return []
+
+    def OPS(self) -> List[OpDef]:
+        """Return the operations defined by this extension"""
+        return [
+            self.state_result_def,
+        ]
+
+    @functools.cached_property
+    def state_result_def(self) -> OpDef:
+        """Report the state of given qubits in the given order.
+
+        This is the generic operation definition. For the instantiated operation, see
+        `stateResult`.
+        """
+        return self().get_op("StateResult")
+
+    @functools.cache
+    def state_result(self, name: str, size: int) -> ExtOp:
+        """Report the state of given qubits in the given order.
+
+        Args:
+            name: The name of the state result to report.
+            size: The size of the qubit array.
+        """
+        return self.state_result_def.instantiate([StringArg(name), BoundedNatArg(size)])

tket-exts/src/tket_exts/tket/futures.py

@@ -0,0 +1,97 @@
+"""Futures extension operations."""
+
+import functools
+from typing import List
+
+from hugr.ext import Extension, OpDef, TypeDef
+from hugr.ops import ExtOp
+from hugr.tys import ExtType, Type, TypeTypeArg
+from ._util import TketExtension, load_extension
+
+
+class FuturesExtension(TketExtension):
+    """Future type and handling operations."""
+
+    @functools.cache
+    def __call__(self) -> Extension:
+        """Returns the futures extension"""
+        return load_extension("tket.futures")
+
+    def TYPES(self) -> List[TypeDef]:
+        """Return the types defined by this extension"""
+        return [self.future_t_def]
+
+    def OPS(self) -> List[OpDef]:
+        """Return the operations defined by this extension"""
+        return [
+            self.dup_def,
+            self.free_def,
+            self.read_def,
+        ]
+
+    @functools.cached_property
+    def future_t_def(self) -> TypeDef:
+        """A value that is computed asynchronously.
+
+        This is the generic type definition. For the instantiated type, see `future_t`
+        """
+        return self().get_type("Future")
+
+    def future_t(self, ty: Type) -> ExtType:
+        """A value that is computed asynchronously.
+
+        Args:
+            The element type wrapped by the Future.
+        """
+        return self().get_type("Future").instantiate([TypeTypeArg(ty)])
+
+    @functools.cached_property
+    def dup_def(self) -> OpDef:
+        """Duplicate a Future. The original is consumed and two Futures are returned.
+
+        This is the generic operation definition. For the instantiated operation, see
+        `dup`.
+        """
+        return self().get_op("Dup")
+
+    def dup(self, ty: Type) -> ExtOp:
+        """Duplicate a Future. The original is consumed and two Futures are returned.
+
+        Args:
+            ty: The element type of the Future being duplicated.
+        """
+        return self.dup_def.instantiate([TypeTypeArg(ty)])
+
+    @functools.cached_property
+    def free_def(self) -> OpDef:
+        """Consume a Future without reading it.
+
+        This is the generic operation definition. For the instantiated operation, see
+        `free`.
+        """
+        return self().get_op("Free")
+
+    def free(self, ty: Type) -> ExtOp:
+        """Consume a Future without reading it.
+
+        Args:
+            ty: The element type of the Future being consumed.
+        """
+        return self.free_def.instantiate([TypeTypeArg(ty)])
+
+    @functools.cached_property
+    def read_def(self) -> OpDef:
+        """Read a value from a Future, consuming it.
+
+        This is the generic operation definition. For the instantiated operation, see
+        `read`.
+        """
+        return self().get_op("Read")
+
+    def read(self, ty: Type) -> ExtOp:
+        """Read a value from a Future, consuming it.
+
+        Args:
+            ty: The element type of the Future being read.
+        """
+        return self.read_def.instantiate([TypeTypeArg(ty)])