guppylang-internals 0.23.0__tar.gz → 0.24.0__tar.gz

{guppylang_internals-0.23.0 → guppylang_internals-0.24.0}/CHANGELOG.md

@@ -3,6 +3,40 @@
 First release of `guppylang_internals` package containing refactored out internal components
 from `guppylang`.
 
+## [0.24.0](https://github.com/CQCL/guppylang/compare/guppylang-internals-v0.23.0...guppylang-internals-v0.24.0) (2025-09-19)
+
+
+### ⚠ BREAKING CHANGES
+
+* `guppylang_internals.decorator.extend_type` now returns a `GuppyDefinition` by default. To get the previous behaviour of returning the annotated class unchanged, pass `return_class=True`.
+* `TypeDef`s now require a `params` field
+* guppylang_internals.ty.parsing.parse_function_io_types replaced with parse_function_arg_annotation and check_function_arg
+* Significant changes to the WASM decorators, types and operations
+* Deleted `guppylang_internals.nodes.{IterHasNext, IterEnd}`
+* guppylang_internals.tracing.unpacking.update_packed_value now returns a bool signalling whether the operation was successful.
+* `CompilationEngine` now initialises all it's fields
+* Calling `CompilationEngine.reset` no longer nullifies `additional_extensions`
+* `CompilationEngine.register_extension` no longer adds duplicates to the `additional_extensions` list
+
+### Features
+
+* Infer type of `self` arguments ([#1192](https://github.com/CQCL/guppylang/issues/1192)) ([51f5a2b](https://github.com/CQCL/guppylang/commit/51f5a2b3a9b06bc4ab054f32a4d07f7395df8ff4))
+
+
+### Bug Fixes
+
+* Add init to CompilationEngine; don't trash additional_extensions ([#1256](https://github.com/CQCL/guppylang/issues/1256)) ([e413748](https://github.com/CQCL/guppylang/commit/e413748532db3895cab4925a222177a4fa3fd61b))
+* Allow generic specialization of methods ([#1206](https://github.com/CQCL/guppylang/issues/1206)) ([93936cc](https://github.com/CQCL/guppylang/commit/93936cc275c56dd856d11fabc7aac20176304147)), closes [#1182](https://github.com/CQCL/guppylang/issues/1182)
+* Correctly update borrowed values after calls and catch cases where it's impossible ([#1253](https://github.com/CQCL/guppylang/issues/1253)) ([3ec5462](https://github.com/CQCL/guppylang/commit/3ec54627729b49689da006a743e9e2c359cd3728))
+* Fix `nat` constructor in comptime functions ([#1258](https://github.com/CQCL/guppylang/issues/1258)) ([e257b6f](https://github.com/CQCL/guppylang/commit/e257b6fc2fe3793d6d8f63feca83bf5ed6643673))
+* Fix incorrect leak error for borrowing functions in comptime ([#1252](https://github.com/CQCL/guppylang/issues/1252)) ([855244e](https://github.com/CQCL/guppylang/commit/855244e2d5e3aeb04c2028f9f2310dba0e74210a)), closes [#1249](https://github.com/CQCL/guppylang/issues/1249)
+* wasm module updates based on tested lowering ([#1230](https://github.com/CQCL/guppylang/issues/1230)) ([657cea2](https://github.com/CQCL/guppylang/commit/657cea27af00a9c02e8d1a3190db535bbd1e7981))
+
+
+### Miscellaneous Chores
+
+* Delete unused old iterator AST nodes ([#1215](https://github.com/CQCL/guppylang/issues/1215)) ([2310897](https://github.com/CQCL/guppylang/commit/231089750e33cf70754e5218feed64053c558c17))
+
 ## [0.23.0](https://github.com/CQCL/guppylang/compare/guppylang-internals-v0.22.0...guppylang-internals-v0.23.0) (2025-08-19)
 
 

{guppylang_internals-0.23.0 → guppylang_internals-0.24.0}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: guppylang-internals
-Version: 0.23.0
+Version: 0.24.0
 Summary: Compiler internals for `guppylang` package.
 Author-email: Mark Koch <mark.koch@quantinuum.com>, TKET development team <tket-support@quantinuum.com>
 Maintainer-email: Mark Koch <mark.koch@quantinuum.com>, TKET development team <tket-support@quantinuum.com>
@@ -220,7 +220,7 @@ Classifier: Programming Language :: Python :: 3.14
 Classifier: Topic :: Software Development :: Compilers
 Requires-Python: <4,>=3.10
 Requires-Dist: hugr~=0.13.1
-Requires-Dist: tket-exts~=0.10.0
+Requires-Dist: tket-exts~=0.11.0
 Requires-Dist: typing-extensions<5,>=4.9.0
 Provides-Extra: pytket
 Requires-Dist: pytket>=1.34; extra == 'pytket'
@@ -228,7 +228,7 @@ Description-Content-Type: text/markdown
 
 # guppylang-internals
 
-This packages contains the internals of the Guppy compiler. 
+This packages contains the internals of the Guppy compiler.
 
 See `guppylang` for the package providing the user-facing language frontend.
 

{guppylang_internals-0.23.0 → guppylang_internals-0.24.0}/README.md

@@ -1,6 +1,6 @@
 # guppylang-internals
 
-This packages contains the internals of the Guppy compiler. 
+This packages contains the internals of the Guppy compiler.
 
 See `guppylang` for the package providing the user-facing language frontend.
 

{guppylang_internals-0.23.0 → guppylang_internals-0.24.0}/pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "guppylang-internals"
-version = "0.23.0"
+version = "0.24.0"
 requires-python = ">=3.10,<4"
 description = "Compiler internals for `guppylang` package."
 license = { file = "LICENCE" }
@@ -34,7 +34,7 @@ classifiers = [
 
 dependencies = [
     "typing-extensions >=4.9.0,<5",
-    "tket-exts ~= 0.10.0",
+    "tket-exts ~= 0.11.0",
     "hugr ~= 0.13.1",
 ]
 

{guppylang_internals-0.23.0 → guppylang_internals-0.24.0}/src/guppylang_internals/__init__.py

@@ -1,3 +1,3 @@
 # This is updated by our release-please workflow, triggered by this
 # annotation: x-release-please-version
-__version__ = "0.23.0"
+__version__ = "0.24.0"

{guppylang_internals-0.23.0 → guppylang_internals-0.24.0}/src/guppylang_internals/checker/core.py

@@ -54,6 +54,7 @@ from guppylang_internals.tys.ty import (
 
 if TYPE_CHECKING:
     from guppylang_internals.definition.struct import StructField
+    from guppylang_internals.tys.parsing import TypeParsingCtx
 
 
 #: A "place" is a description for a storage location of a local value that users
@@ -507,6 +508,13 @@ class Context(NamedTuple):
     locals: Locals[str, Variable]
     generic_params: dict[str, Parameter]
 
+    @property
+    def parsing_ctx(self) -> "TypeParsingCtx":
+        """A type parsing context derived from this checking context."""
+        from guppylang_internals.tys.parsing import TypeParsingCtx
+
+        return TypeParsingCtx(self.globals, self.generic_params)
+
 
 class DummyEvalDict(dict[str, Any]):
     """A custom dict that can be passed to `eval` to give better error messages.

{guppylang_internals-0.23.0 → guppylang_internals-0.24.0}/src/guppylang_internals/checker/expr_checker.py

@@ -34,6 +34,7 @@ from guppylang_internals.ast_util import (
     AstNode,
     AstVisitor,
     breaks_in_loop,
+    get_type,
     get_type_opt,
     return_nodes_in_ast,
     with_loc,
@@ -101,8 +102,6 @@ from guppylang_internals.nodes import (
     FieldAccessAndDrop,
     GenericParamValue,
     GlobalName,
-    IterEnd,
-    IterHasNext,
     IterNext,
     LocalCall,
     MakeIter,
@@ -784,14 +783,6 @@ class ExprSynthesizer(AstVisitor[tuple[ast.expr, Type]]):
                 raise GuppyTypeError(err)
         return expr, ty
 
-    def visit_IterHasNext(self, node: IterHasNext) -> tuple[ast.expr, Type]:
-        node.value, ty = self.synthesize(node.value)
-        flags = InputFlags.Owned if not ty.copyable else InputFlags.NoFlags
-        exp_sig = FunctionType([FuncInput(ty, flags)], TupleType([bool_type(), ty]))
-        return self.synthesize_instance_func(
-            node.value, [], "__hasnext__", "an iterator", exp_sig, True
-        )
-
     def visit_IterNext(self, node: IterNext) -> tuple[ast.expr, Type]:
         node.value, ty = self.synthesize(node.value)
         flags = InputFlags.Owned if not ty.copyable else InputFlags.NoFlags
@@ -803,14 +794,6 @@ class ExprSynthesizer(AstVisitor[tuple[ast.expr, Type]]):
             node.value, [], "__next__", "an iterator", exp_sig, True
         )
 
-    def visit_IterEnd(self, node: IterEnd) -> tuple[ast.expr, Type]:
-        node.value, ty = self.synthesize(node.value)
-        flags = InputFlags.Owned if not ty.copyable else InputFlags.NoFlags
-        exp_sig = FunctionType([FuncInput(ty, flags)], NoneType())
-        return self.synthesize_instance_func(
-            node.value, [], "__end__", "an iterator", exp_sig, True
-        )
-
     def visit_ListComp(self, node: ast.ListComp) -> tuple[ast.expr, Type]:
         raise InternalGuppyError(
             "BB contains `ListComp`. Should have been removed during CFG"
@@ -946,7 +929,7 @@ def check_type_apply(ty: FunctionType, node: ast.Subscript, ctx: Context) -> Ins
         raise GuppyError(err)
 
     return [
-        param.check_arg(arg_from_ast(arg_expr, globals, ctx.generic_params), arg_expr)
+        param.check_arg(arg_from_ast(arg_expr, ctx.parsing_ctx), arg_expr)
         for arg_expr, param in zip(arg_exprs, ty.params, strict=True)
     ]
 
@@ -1232,7 +1215,14 @@ def instantiate_poly(node: ast.expr, ty: FunctionType, inst: Inst) -> ast.expr:
     """Instantiates quantified type arguments in a function."""
     assert len(ty.params) == len(inst)
     if len(inst) > 0:
-        node = with_loc(node, TypeApply(value=with_type(ty, node), inst=inst))
+        # Partial applications need to be instantiated on the inside
+        if isinstance(node, PartialApply):
+            full_ty = get_type(node.func)
+            assert isinstance(full_ty, FunctionType)
+            assert full_ty.params == ty.params
+            node.func = instantiate_poly(node.func, full_ty, inst)
+        else:
+            node = with_loc(node, TypeApply(value=with_type(ty, node), inst=inst))
         return with_type(ty.instantiate(inst), node)
     return with_type(ty, node)
 

{guppylang_internals-0.23.0 → guppylang_internals-0.24.0}/src/guppylang_internals/checker/func_checker.py

@@ -7,8 +7,8 @@ node straight from the Python AST. We build a CFG, check it, and return a
 
 import ast
 import sys
-from dataclasses import dataclass
-from typing import TYPE_CHECKING, ClassVar
+from dataclasses import dataclass, replace
+from typing import TYPE_CHECKING, ClassVar, cast
 
 from guppylang_internals.ast_util import return_nodes_in_ast, with_loc
 from guppylang_internals.cfg.bb import BB
@@ -17,13 +17,28 @@ from guppylang_internals.checker.cfg_checker import CheckedCFG, check_cfg
 from guppylang_internals.checker.core import Context, Globals, Place, Variable
 from guppylang_internals.checker.errors.generic import UnsupportedError
 from guppylang_internals.definition.common import DefId
+from guppylang_internals.definition.ty import TypeDef
 from guppylang_internals.diagnostic import Error, Help, Note
 from guppylang_internals.engine import DEF_STORE, ENGINE
 from guppylang_internals.error import GuppyError
 from guppylang_internals.experimental import check_capturing_closures_enabled
 from guppylang_internals.nodes import CheckedNestedFunctionDef, NestedFunctionDef
-from guppylang_internals.tys.parsing import parse_function_io_types
-from guppylang_internals.tys.ty import FunctionType, InputFlags, NoneType
+from guppylang_internals.tys.parsing import (
+    TypeParsingCtx,
+    check_function_arg,
+    parse_function_arg_annotation,
+    type_from_ast,
+    type_with_flags_from_ast,
+)
+from guppylang_internals.tys.ty import (
+    ExistentialTypeVar,
+    FuncInput,
+    FunctionType,
+    InputFlags,
+    NoneType,
+    Type,
+    unify,
+)
 
 if sys.version_info >= (3, 12):
     from guppylang_internals.tys.parsing import parse_parameter
@@ -53,6 +68,15 @@ class MissingArgAnnotationError(Error):
     span_label: ClassVar[str] = "Argument requires a type annotation"
 
 
+@dataclass(frozen=True)
+class RecursiveSelfError(Error):
+    title: ClassVar[str] = "Recursive self annotation"
+    span_label: ClassVar[str] = (
+        "Type of `{self_arg}` cannot recursively refer to `Self`"
+    )
+    self_arg: str
+
+
 @dataclass(frozen=True)
 class MissingReturnAnnotationError(Error):
     title: ClassVar[str] = "Missing type annotation"
@@ -67,6 +91,43 @@ class MissingReturnAnnotationError(Error):
         func: str
 
 
+@dataclass(frozen=True)
+class InvalidSelfError(Error):
+    title: ClassVar[str] = "Invalid self annotation"
+    span_label: ClassVar[str] = "`{self_arg}` must be of type `{self_ty}`"
+    self_arg: str
+    self_ty: Type
+
+
+@dataclass(frozen=True)
+class SelfParamsShadowedError(Error):
+    title: ClassVar[str] = "Shadowed generic parameters"
+    span_label: ClassVar[str] = (
+        "Cannot infer type for `{self_arg}` since parameter `{param}` of "
+        "`{ty_defn.name}` is shadowed"
+    )
+    param: str
+    ty_defn: "TypeDef"
+    self_arg: str
+
+    @dataclass(frozen=True)
+    class ExplicitHelp(Help):
+        span_label: ClassVar[str] = (
+            "Consider specifying the type explicitly: `{suggestion}`"
+        )
+
+        @property
+        def suggestion(self) -> str:
+            parent = self._parent
+            assert isinstance(parent, SelfParamsShadowedError)
+            params = (
+                f"[{', '.join(f'?{p.name}' for p in parent.ty_defn.params)}]"
+                if parent.ty_defn.params
+                else ""
+            )
+            return f'{parent.self_arg}: "{parent.ty_defn.name}{params}"'
+
+
 def check_global_func_def(
     func_def: ast.FunctionDef, ty: FunctionType, globals: Globals
 ) -> CheckedCFG[Place]:
@@ -176,9 +237,16 @@ def check_nested_func_def(
     return with_loc(func_def, checked_def)
 
 
-def check_signature(func_def: ast.FunctionDef, globals: Globals) -> FunctionType:
+def check_signature(
+    func_def: ast.FunctionDef, globals: Globals, def_id: DefId | None = None
+) -> FunctionType:
     """Checks the signature of a function definition and returns the corresponding
-    Guppy type."""
+    Guppy type.
+
+    If this is a method, then the `DefId` of the associated parent type should also be
+    passed. This will be used to check or infer the type annotation for the `self`
+    argument.
+    """
     if len(func_def.args.posonlyargs) != 0:
         raise GuppyError(
             UnsupportedError(func_def.args.posonlyargs[0], "Positional-only parameters")
@@ -211,23 +279,29 @@ def check_signature(func_def: ast.FunctionDef, globals: Globals) -> FunctionType
             param = parse_parameter(param_node, i, globals)
             param_var_mapping[param.name] = param
 
-    input_nodes = []
+    # Figure out if this is a method
+    self_defn: TypeDef | None = None
+    if def_id is not None and def_id in DEF_STORE.impl_parents:
+        self_defn = cast(TypeDef, ENGINE.get_checked(DEF_STORE.impl_parents[def_id]))
+        assert isinstance(self_defn, TypeDef)
+
+    inputs = []
     input_names = []
-    for inp in func_def.args.args:
-        ty_ast = inp.annotation
-        if ty_ast is None:
-            raise GuppyError(MissingArgAnnotationError(inp))
-        input_nodes.append(ty_ast)
+    ctx = TypeParsingCtx(globals, param_var_mapping, allow_free_vars=True)
+    for i, inp in enumerate(func_def.args.args):
+        # Special handling for `self` arguments. Note that `__new__` is excluded here
+        # since it's not a method so doesn't take `self`.
+        if self_defn and i == 0 and func_def.name != "__new__":
+            input = parse_self_arg(inp, self_defn, ctx)
+            ctx = replace(ctx, self_ty=input.ty)
+        else:
+            ty_ast = inp.annotation
+            if ty_ast is None:
+                raise GuppyError(MissingArgAnnotationError(inp))
+            input = parse_function_arg_annotation(ty_ast, inp.arg, ctx)
+        inputs.append(input)
         input_names.append(inp.arg)
-    inputs, output = parse_function_io_types(
-        input_nodes,
-        func_def.returns,
-        input_names,
-        func_def,
-        globals,
-        param_var_mapping,
-        True,
-    )
+    output = type_from_ast(func_def.returns, ctx)
     return FunctionType(
         inputs,
         output,
@@ -236,6 +310,81 @@ def check_signature(func_def: ast.FunctionDef, globals: Globals) -> FunctionType
     )
 
 
+def parse_self_arg(arg: ast.arg, self_defn: TypeDef, ctx: TypeParsingCtx) -> FuncInput:
+    """Handles parsing of the `self` argument on methods.
+
+    This argument is special since its type annotation may be omitted. Furthermore, if a
+    type is provided then it must match the parent type.
+    """
+    assert self_defn.params is not None
+    if arg.annotation is None:
+        return handle_implicit_self_arg(arg, self_defn, ctx)
+
+    # If the user has provided an annotation for `self`, then we go ahead and parse it.
+    # However, in the annotation the user is also allowed to use `Self`, so we have to
+    # specify a `self_ty` in the context.
+    self_ty_head = self_defn.check_instantiate(
+        [param.to_existential()[0] for param in self_defn.params]
+    )
+    self_ty_placeholder = ExistentialTypeVar.fresh(
+        "Self", copyable=self_ty_head.copyable, droppable=self_ty_head.droppable
+    )
+    assert ctx.self_ty is None
+    ctx = replace(ctx, self_ty=self_ty_placeholder)
+    user_ty, user_flags = type_with_flags_from_ast(arg.annotation, ctx)
+
+    # If the user just annotates `self: Self` then we can fall back to the case where
+    # no annotation is provided at all
+    if user_ty == self_ty_placeholder:
+        return handle_implicit_self_arg(arg, self_defn, ctx, user_flags)
+
+    # Annotations like `self: Foo[Self]` are not allowed (would be an infinite type)
+    if self_ty_placeholder in user_ty.unsolved_vars:
+        raise GuppyError(RecursiveSelfError(arg.annotation, arg.arg))
+
+    # Check that the annotation matches the parent type. We can do this by unifying with
+    # the expected self type where all params are instantiated with unification vars
+    subst = unify(user_ty, self_ty_head, {})
+    if subst is None:
+        raise GuppyError(InvalidSelfError(arg.annotation, arg.arg, self_ty_head))
+
+    return check_function_arg(user_ty, user_flags, arg, arg.arg, ctx)
+
+
+def handle_implicit_self_arg(
+    arg: ast.arg,
+    self_defn: TypeDef,
+    ctx: TypeParsingCtx,
+    flags: InputFlags = InputFlags.NoFlags,
+) -> FuncInput:
+    """Handles the case where no annotation for `self` is provided.
+
+    Generates the most generic annotation that is possible by making the function as
+    generic as the parent type.
+    """
+    # Check that the user hasn't shadowed some of the parent type parameters using a
+    # Python 3.12 style parameter declaration
+    assert self_defn.params is not None
+    shadowed_params = [
+        param for param in self_defn.params if param.name in ctx.param_var_mapping
+    ]
+    if shadowed_params:
+        param = shadowed_params.pop()
+        err = SelfParamsShadowedError(arg, param.name, self_defn, arg.arg)
+        err.add_sub_diagnostic(SelfParamsShadowedError.ExplicitHelp(arg))
+        raise GuppyError(err)
+
+    # The generic params inherited from the parent type should appear first in the
+    # parameter list, so we have to shift the existing ones
+    for name, param in ctx.param_var_mapping.items():
+        ctx.param_var_mapping[name] = param.with_idx(param.idx + len(self_defn.params))
+
+    ctx.param_var_mapping.update({param.name: param for param in self_defn.params})
+    self_args = [param.to_bound() for param in self_defn.params]
+    self_ty = self_defn.check_instantiate(self_args, loc=arg)
+    return check_function_arg(self_ty, flags, arg, arg.arg, ctx)
+
+
 def parse_function_with_docstring(
     func_ast: ast.FunctionDef,
 ) -> tuple[ast.FunctionDef, str | None]:

{guppylang_internals-0.23.0 → guppylang_internals-0.24.0}/src/guppylang_internals/checker/stmt_checker.py

@@ -356,7 +356,7 @@ class StmtChecker(AstVisitor[BBStatement]):
     def visit_AnnAssign(self, node: ast.AnnAssign) -> ast.stmt:
         if node.value is None:
             raise GuppyError(UnsupportedError(node, "Variable declarations"))
-        ty = type_from_ast(node.annotation, self.ctx.globals, self.ctx.generic_params)
+        ty = type_from_ast(node.annotation, self.ctx.parsing_ctx)
         node.value, subst = self._check_expr(node.value, ty)
         assert not ty.unsolved_vars  # `ty` must be closed!
         assert len(subst) == 0

{guppylang_internals-0.23.0 → guppylang_internals-0.24.0}/src/guppylang_internals/decorator.py

@@ -1,7 +1,7 @@
 from __future__ import annotations
 
 import inspect
-from typing import TYPE_CHECKING, ParamSpec, TypeVar
+from typing import TYPE_CHECKING, ParamSpec, TypeVar, overload
 
 from hugr import ops
 from hugr import tys as ht
@@ -42,6 +42,7 @@ from guppylang_internals.tys.ty import (
 )
 
 if TYPE_CHECKING:
+    import ast
     import builtins
     from collections.abc import Callable, Sequence
     from types import FrameType
@@ -121,15 +122,19 @@ def hugr_op(
     return custom_function(OpCompiler(op), checker, higher_order_value, name, signature)
 
 
-def extend_type(defn: TypeDef) -> Callable[[type], type]:
-    """Decorator to add new instance functions to a type."""
+def extend_type(defn: TypeDef, return_class: bool = False) -> Callable[[type], type]:
+    """Decorator to add new instance functions to a type.
+
+    By default, returns a `GuppyDefinition` object referring to the type. Alternatively,
+    `return_class=True` can be set to return the decorated class unchanged.
+    """
     from guppylang.defs import GuppyDefinition
 
     def dec(c: type) -> type:
         for val in c.__dict__.values():
             if isinstance(val, GuppyDefinition):
                 DEF_STORE.register_impl(defn.id, val.wrapped.name, val.id)
-        return c
+        return c if return_class else GuppyDefinition(defn)  # type: ignore[return-value]
 
     return dec
 
@@ -181,63 +186,124 @@ def custom_type(
 
 
 def wasm_module(
-    filename: str, filehash: int
+    filename: str,
 ) -> Callable[[builtins.type[T]], GuppyDefinition]:
-    from guppylang.defs import GuppyDefinition
-
-    def dec(cls: builtins.type[T]) -> GuppyDefinition:
-        # N.B. Only one module per file and vice-versa
-        wasm_module = WasmModuleTypeDef(
-            DefId.fresh(),
-            cls.__name__,
-            None,
-            filename,
-            filehash,
-        )
-
-        wasm_module_ty = wasm_module.check_instantiate([], None)
-
-        DEF_STORE.register_def(wasm_module, get_calling_frame())
-        for val in cls.__dict__.values():
-            if isinstance(val, GuppyDefinition):
-                DEF_STORE.register_impl(wasm_module.id, val.wrapped.name, val.id)
-        # Add a constructor to the class
-        call_method = CustomFunctionDef(
-            DefId.fresh(),
-            "__new__",
-            None,
-            FunctionType(
-                [FuncInput(NumericType(NumericType.Kind.Nat), flags=InputFlags.Owned)],
-                wasm_module_ty,
-            ),
-            DefaultCallChecker(),
-            WasmModuleInitCompiler(),
-            True,
-            GlobalConstId.fresh(f"{cls.__name__}.__new__"),
-            True,
-        )
-        discard = CustomFunctionDef(
-            DefId.fresh(),
-            "discard",
-            None,
-            FunctionType([FuncInput(wasm_module_ty, InputFlags.Owned)], NoneType()),
-            DefaultCallChecker(),
-            WasmModuleDiscardCompiler(),
-            False,
-            GlobalConstId.fresh(f"{cls.__name__}.__discard__"),
-            True,
-        )
-        DEF_STORE.register_def(call_method, get_calling_frame())
-        DEF_STORE.register_impl(wasm_module.id, "__new__", call_method.id)
-        DEF_STORE.register_def(discard, get_calling_frame())
-        DEF_STORE.register_impl(wasm_module.id, "discard", discard.id)
+    def type_def_wrapper(
+        id: DefId,
+        name: str,
+        defined_at: ast.AST | None,
+        wasm_file: str,
+        config: str | None,
+    ) -> OpaqueTypeDef:
+        assert config is None
+        return WasmModuleTypeDef(id, name, defined_at, wasm_file)
+
+    f = ext_module_decorator(
+        type_def_wrapper, WasmModuleInitCompiler(), WasmModuleDiscardCompiler(), True
+    )
+    return f(filename, None)
 
-        return GuppyDefinition(wasm_module)
-
-    return dec
 
+def ext_module_decorator(
+    type_def: Callable[[DefId, str, ast.AST | None, str, str | None], OpaqueTypeDef],
+    init_compiler: CustomInoutCallCompiler,
+    discard_compiler: CustomInoutCallCompiler,
+    init_arg: bool,  # Whether the init function should take a nat argument
+) -> Callable[[str, str | None], Callable[[builtins.type[T]], GuppyDefinition]]:
+    from guppylang.defs import GuppyDefinition
 
-def wasm(f: Callable[P, T]) -> GuppyFunctionDefinition[P, T]:
+    def fun(
+        filename: str, module: str | None
+    ) -> Callable[[builtins.type[T]], GuppyDefinition]:
+        def dec(cls: builtins.type[T]) -> GuppyDefinition:
+            # N.B. Only one module per file and vice-versa
+            ext_module = type_def(
+                DefId.fresh(),
+                cls.__name__,
+                None,
+                filename,
+                module,
+            )
+
+            ext_module_ty = ext_module.check_instantiate([], None)
+
+            DEF_STORE.register_def(ext_module, get_calling_frame())
+            for val in cls.__dict__.values():
+                if isinstance(val, GuppyDefinition):
+                    DEF_STORE.register_impl(ext_module.id, val.wrapped.name, val.id)
+            # Add a constructor to the class
+            if init_arg:
+                init_fn_ty = FunctionType(
+                    [
+                        FuncInput(
+                            NumericType(NumericType.Kind.Nat),
+                            flags=InputFlags.Owned,
+                        )
+                    ],
+                    ext_module_ty,
+                )
+            else:
+                init_fn_ty = FunctionType([], ext_module_ty)
+
+            call_method = CustomFunctionDef(
+                DefId.fresh(),
+                "__new__",
+                None,
+                init_fn_ty,
+                DefaultCallChecker(),
+                init_compiler,
+                True,
+                GlobalConstId.fresh(f"{cls.__name__}.__new__"),
+                True,
+            )
+            discard = CustomFunctionDef(
+                DefId.fresh(),
+                "discard",
+                None,
+                FunctionType([FuncInput(ext_module_ty, InputFlags.Owned)], NoneType()),
+                DefaultCallChecker(),
+                discard_compiler,
+                False,
+                GlobalConstId.fresh(f"{cls.__name__}.__discard__"),
+                True,
+            )
+            DEF_STORE.register_def(call_method, get_calling_frame())
+            DEF_STORE.register_impl(ext_module.id, "__new__", call_method.id)
+            DEF_STORE.register_def(discard, get_calling_frame())
+            DEF_STORE.register_impl(ext_module.id, "discard", discard.id)
+
+            return GuppyDefinition(ext_module)
+
+        return dec
+
+    return fun
+
+
+@overload
+def wasm(arg: Callable[P, T]) -> GuppyFunctionDefinition[P, T]: ...
+
+
+@overload
+def wasm(arg: int) -> Callable[[Callable[P, T]], GuppyFunctionDefinition[P, T]]: ...
+
+
+def wasm(
+    arg: int | Callable[P, T],
+) -> (
+    GuppyFunctionDefinition[P, T]
+    | Callable[[Callable[P, T]], GuppyFunctionDefinition[P, T]]
+):
+    if isinstance(arg, int):
+
+        def wrapper(f: Callable[P, T]) -> GuppyFunctionDefinition[P, T]:
+            return wasm_helper(arg, f)
+
+        return wrapper
+    else:
+        return wasm_helper(None, arg)
+
+
+def wasm_helper(fn_id: int | None, f: Callable[P, T]) -> GuppyFunctionDefinition[P, T]:
     from guppylang.defs import GuppyFunctionDefinition
 
     func = RawWasmFunctionDef(
@@ -246,7 +312,7 @@ def wasm(f: Callable[P, T]) -> GuppyFunctionDefinition[P, T]:
         None,
         f,
         WasmCallChecker(),
-        WasmModuleCallCompiler(f.__name__),
+        WasmModuleCallCompiler(f.__name__, fn_id),
         True,
         signature=None,
     )