fix: Ignore docstrings in guppy.unitary classes

guppylang-internals/src/guppylang_internals/checker/unitary_checker.py

@@ -611,44 +611,61 @@ def check_unitary_method[T](
 
     custom_methods: dict[CustomModifierKind, RawFunctionDef] = {}
     valid_method_names = [kind.value for kind in CustomModifierKind] + ["__call__"]
-    for node in class_ast.body:
-        if not isinstance(node, ast.FunctionDef) or node.name not in valid_method_names:
-            raise GuppyError(
-                UnexpectedError(
-                    node, "statement", unexpected_in="in a `@guppy.unitary` class"
-                ).add_sub_diagnostic(InvalidUnitaryMethodHelp(None))
-            )
-        method_name = node.name
-        # `__call__` has already been checked in `decorator._get_unitary_call_def`
-        if method_name == "__call__":
-            continue
-        method = cls.__dict__.get(method_name)
-
-        if not (
-            isinstance(method, GuppyDefinition)
-            and isinstance(method.wrapped, RawFunctionDef)
-        ):
-            raise GuppyError(InvalidUnitaryMethodError(node, node.name, class_ast.name))
+    for i, node in enumerate(class_ast.body):
+        match i, node:
+            # Docstrings are fine if they occur at the start
+            case 0, ast.Expr(value=ast.Constant(value=v)) if isinstance(v, str):
+                pass
+            case _, ast.FunctionDef() as node:
+                if node.name not in valid_method_names:
+                    raise GuppyError(
+                        UnexpectedError(
+                            node,
+                            "statement",
+                            unexpected_in="in a `@guppy.unitary` class",
+                        ).add_sub_diagnostic(InvalidUnitaryMethodHelp(None))
+                    )
+                method_name = node.name
+                # `__call__` already checked in `decorator._get_unitary_call_def`
+                if method_name == "__call__":
+                    continue
+                method = cls.__dict__.get(method_name)
+
+                if not (
+                    isinstance(method, GuppyDefinition)
+                    and isinstance(method.wrapped, RawFunctionDef)
+                ):
+                    raise GuppyError(
+                        InvalidUnitaryMethodError(node, node.name, class_ast.name)
+                    )
 
-        # Check that no invalid metadata is present
-        method_raw_def = method.wrapped
-        if method_raw_def.unitary_flags != UnitaryFlags.NoFlags:
-            raise GuppyError(
-                UnexpectedError(
-                    function_header_span(node),
-                    "unitary flags",
-                    unexpected_in="in the method `@guppy` decorator",
-                ).add_sub_diagnostic(InvalidUnitaryMetadataHelp(None))
-            )
+                # Check that no invalid metadata is present
+                method_raw_def = method.wrapped
+                if method_raw_def.unitary_flags != UnitaryFlags.NoFlags:
+                    raise GuppyError(
+                        UnexpectedError(
+                            function_header_span(node),
+                            "unitary flags",
+                            unexpected_in="in the method `@guppy` decorator",
+                        ).add_sub_diagnostic(InvalidUnitaryMetadataHelp(None))
+                    )
 
-        if (
-            method_raw_def.metadata is not None
-            and method_raw_def.metadata.get_expected_qubits() is not None
-        ):
-            raise GuppyError(
-                UnexpectedError(
-                    function_header_span(node), "`@expected_qubits` decorator"
-                ).add_sub_diagnostic(InvalidExpectedQubit(None))
-            )
-        custom_methods[CustomModifierKind(node.name)] = method_raw_def
+                if (
+                    method_raw_def.metadata is not None
+                    and method_raw_def.metadata.get_expected_qubits() is not None
+                ):
+                    raise GuppyError(
+                        UnexpectedError(
+                            function_header_span(node), "`@expected_qubits` decorator"
+                        ).add_sub_diagnostic(InvalidExpectedQubit(None))
+                    )
+                custom_methods[CustomModifierKind(node.name)] = method_raw_def
+            case _, _:
+                raise GuppyError(
+                    UnexpectedError(
+                        node,
+                        "statement",
+                        unexpected_in="in a `@guppy.unitary` class",
+                    ).add_sub_diagnostic(InvalidUnitaryMethodHelp(None))
+                )
     return custom_methods

tests/integration/test_modifier.py

@@ -1698,3 +1698,28 @@ def main(
         ps[i].b = ps[i].a
 
     validate(main.with_minimal_opt().compile_function(), export=False)
+
+
+def test_custom_modifier_accepts_docstring(validate) -> None:
+    """guppy.unitary classes accept docstrings."""
+
+    @guppy.unitary
+    class custom_modified:
+        """Docstring here"""
+
+        @guppy
+        def __call__(q: qubit) -> None:
+            pass
+
+        @guppy
+        def daggered(q: qubit) -> None:
+            pass
+
+    @guppy
+    def main() -> None:
+        q = qubit()
+        with dagger:
+            custom_modified(q)
+        discard(q)
+
+    validate(main.compile())