[4/5] Introduce OperationDefinition (#16)

Author: dandavison

src/nexusrpc/__init__.py

@@ -28,16 +28,16 @@
     OutputT,
 )
 from ._serializer import Content, LazyValue
-from ._service import Operation, ServiceDefinition, service
+from ._service import Operation, OperationDefinition, ServiceDefinition, service
 from ._util import (
-    get_operation_definition,
+    get_operation,
     get_service_definition,
-    set_operation_definition,
+    set_operation,
 )
 
 __all__ = [
     "Content",
-    "get_operation_definition",
+    "get_operation",
     "get_service_definition",
     "handler",
     "HandlerError",
@@ -46,12 +46,13 @@
     "LazyValue",
     "Link",
     "Operation",
+    "OperationDefinition",
     "OperationError",
     "OperationErrorState",
     "OperationInfo",
     "OperationState",
     "OutputT",
     "service",
     "ServiceDefinition",
-    "set_operation_definition",
+    "set_operation",
 ]

src/nexusrpc/_service.py

@@ -51,23 +51,38 @@ class MyNexusService:
     input_type: Optional[type[InputT]] = dataclasses.field(default=None)
     output_type: Optional[type[OutputT]] = dataclasses.field(default=None)
 
-    def __post_init__(self):
-        if not self.name:
-            raise ValueError("Operation name cannot be empty")
 
-    def _validation_errors(self) -> list[str]:
-        errors = []
-        if not self.name:
-            errors.append(
-                f"Operation has no name (method_name is '{self.method_name}')"
+@dataclass
+class OperationDefinition(Generic[InputT, OutputT]):
+    """
+    Internal representation of a user's :py:class:`Operation` definition.
+    """
+
+    name: str
+    method_name: str
+    input_type: type[InputT]
+    output_type: type[OutputT]
+
+    @classmethod
+    def from_operation(
+        cls, operation: Operation[InputT, OutputT]
+    ) -> OperationDefinition[InputT, OutputT]:
+        if not operation.name:
+            raise ValueError(
+                f"Operation has no name (method_name is '{operation.method_name}')"
             )
-        if not self.method_name:
-            errors.append(f"Operation '{self.name}' has no method name")
-        if not self.input_type:
-            errors.append(f"Operation '{self.name}' has no input type")
-        if not self.output_type:
-            errors.append(f"Operation '{self.name}' has no output type")
-        return errors
+        if not operation.method_name:
+            raise ValueError(f"Operation '{operation.name}' has no method name")
+        if not operation.input_type:
+            raise ValueError(f"Operation '{operation.name}' has no input type")
+        if not operation.output_type:
+            raise ValueError(f"Operation '{operation.name}' has no output type")
+        return cls(
+            name=operation.name,
+            method_name=operation.method_name,
+            input_type=operation.input_type,
+            output_type=operation.output_type,
+        )
 
 
 @overload
@@ -124,16 +139,23 @@ def decorator(cls: type[ServiceT]) -> type[ServiceT]:
         defn = ServiceDefinition.from_class(cls, name or cls.__name__)
         set_service_definition(cls, defn)
 
-        # In order for callers to refer to operations at run-time, a decorated user
+        # In order for callers to refer to operation definitions at run-time, a decorated user
         # service class must itself have a class attribute for every operation, even if
         # declared only via a type annotation, and whether inherited from a parent class
         # or not.
         #
         # TODO(preview): it is sufficient to do this setattr only for the subset of
         # operations that were declared on *this* class. Currently however we are
         # setting all inherited operations.
-        for op_name, op in defn.operations.items():
-            setattr(cls, op_name, op)
+        for op_name, op_defn in defn.operation_definitions.items():
+            if not hasattr(cls, op_name):
+                op = Operation(
+                    name=op_defn.name,
+                    method_name=op_defn.method_name,
+                    input_type=op_defn.input_type,
+                    output_type=op_defn.output_type,
+                )
+                setattr(cls, op_name, op)
 
         return cls
 
@@ -155,7 +177,7 @@ class ServiceDefinition:
     """
 
     name: str
-    operations: Mapping[str, Operation[Any, Any]]
+    operation_definitions: Mapping[str, OperationDefinition[Any, Any]]
 
     def __post_init__(self):
         if errors := self._validation_errors():
@@ -173,7 +195,7 @@ def from_class(user_class: type[ServiceT], name: str) -> ServiceDefinition:
         If multiple service definitions define an operation with the same name, then the
         usual mro() precedence rules apply.
         """
-        operations = ServiceDefinition._collect_operations(user_class)
+        operation_definitions = ServiceDefinition._collect_operations(user_class)
 
         # Obtain the set of operations to be inherited from ancestral service
         # definitions. Operations are only inherited from classes that are also
@@ -186,47 +208,53 @@ def from_class(user_class: type[ServiceT], name: str) -> ServiceDefinition:
         # 2. No inherited operation has the same method name as that of an operation
         #    defined here. If this were violated, there would be ambiguity in which
         #    operation handler is dispatched to.
-        parent_defns = (
-            defn
-            for defn in (get_service_definition(cls) for cls in user_class.mro()[1:])
-            if defn
+        parent_service_defn = next(
+            (
+                defn
+                for defn in (
+                    get_service_definition(cls) for cls in user_class.mro()[1:]
+                )
+                if defn
+            ),
+            None,
         )
-        method_names = {op.method_name for op in operations.values() if op.method_name}
-        if parent_defn := next(parent_defns, None):
-            for op in parent_defn.operations.values():
-                if op.method_name in method_names:
+        if parent_service_defn:
+            method_names = {op.method_name for op in operation_definitions.values()}
+            for op_defn in parent_service_defn.operation_definitions.values():
+                if op_defn.method_name in method_names:
                     raise ValueError(
-                        f"Operation method name '{op.method_name}' in class '{user_class}' "
+                        f"Operation method name '{op_defn.method_name}' in class '{user_class}' "
                         f"also occurs in a service definition inherited from a parent class: "
-                        f"'{parent_defn.name}'. This is not allowed."
+                        f"'{parent_service_defn.name}'. This is not allowed."
                     )
-                if op.name in operations:
+                if op_defn.name in operation_definitions:
                     raise ValueError(
-                        f"Operation name '{op.name}' in class '{user_class}' "
+                        f"Operation name '{op_defn.name}' in class '{user_class}' "
                         f"also occurs in a service definition inherited from a parent class: "
-                        f"'{parent_defn.name}'. This is not allowed."
+                        f"'{parent_service_defn.name}'. This is not allowed."
                     )
-                operations[op.name] = op
+                operation_definitions[op_defn.name] = op_defn
 
-        return ServiceDefinition(name=name, operations=operations)
+        return ServiceDefinition(name=name, operation_definitions=operation_definitions)
 
     def _validation_errors(self) -> list[str]:
         errors = []
         if not self.name:
             errors.append("Service has no name")
         seen_method_names = set()
-        for op in self.operations.values():
-            if op.method_name in seen_method_names:
-                errors.append(f"Operation method name '{op.method_name}' is not unique")
-            seen_method_names.add(op.method_name)
-            errors.extend(op._validation_errors())  # pyright: ignore[reportPrivateUsage]
+        for op_defn in self.operation_definitions.values():
+            if op_defn.method_name in seen_method_names:
+                errors.append(
+                    f"Operation method name '{op_defn.method_name}' is not unique"
+                )
+            seen_method_names.add(op_defn.method_name)
         return errors
 
     @staticmethod
     def _collect_operations(
         user_class: type[ServiceT],
-    ) -> dict[str, Operation[Any, Any]]:
-        """Collect operations from a user service definition class.
+    ) -> dict[str, OperationDefinition[Any, Any]]:
+        """Collect operation definitions from a user service definition class.
 
         Does not visit parent classes.
         """
@@ -301,11 +329,11 @@ def _collect_operations(
             if op.method_name is None:
                 op.method_name = key
 
-        operations_by_name = {}
+        op_defns = {}
         for op in operations.values():
-            if op.name in operations_by_name:
-                raise ValueError(
-                    f"Operation '{op.name}' in class '{user_class}' is defined multiple times"
+            if op.name in op_defns:
+                raise RuntimeError(
+                    f"Operation '{op.name}' in service '{user_class}' is defined multiple times"
                 )
-            operations_by_name[op.name] = op
-        return operations_by_name
+            op_defns[op.name] = OperationDefinition.from_operation(op)
+        return op_defns

src/nexusrpc/_util.py

@@ -16,20 +16,25 @@
     from nexusrpc._common import ServiceT
     from nexusrpc.handler._operation_handler import OperationHandler
 
+_NEXUS_SERVICE_DEFINITION_ATTR_NAME = "__nexus_service_definition__"
+_NEXUS_OPERATION_ATTR_NAME = "__nexus_operation__"
+_NEXUS_OPERATION_FACTORY_ATTR_NAME = "__nexus_operation_factory__"
+
 
 def get_service_definition(
     obj: Any,
 ) -> Optional[nexusrpc.ServiceDefinition]:
     """Return the :py:class:`nexusrpc.ServiceDefinition` for the object, or None"""
-    # getattr would allow a non-decorated class to act as a service
+    # Do not use getattr since it would allow a non-decorated class to act as a service
     # definition if it inherits from a decorated class.
     if isinstance(obj, type):
-        defn = obj.__dict__.get("__nexus_service__")
+        defn = obj.__dict__.get(_NEXUS_SERVICE_DEFINITION_ATTR_NAME)
     else:
-        defn = getattr(obj, "__dict__", {}).get("__nexus_service__")
+        defn = getattr(obj, "__dict__", {}).get(_NEXUS_SERVICE_DEFINITION_ATTR_NAME)
     if defn and not isinstance(defn, nexusrpc.ServiceDefinition):
         raise ValueError(
-            f"Service definition {obj.__name__} has a __nexus_service__ attribute that is not a ServiceDefinition."
+            f"{obj.__name__} has a {_NEXUS_SERVICE_DEFINITION_ATTR_NAME} attribute "
+            f"that is not a nexusrpc.ServiceDefinition."
         )
     return defn
 
@@ -38,28 +43,33 @@ def set_service_definition(
     cls: type[ServiceT], service_definition: nexusrpc.ServiceDefinition
 ) -> None:
     """Set the :py:class:`nexusrpc.ServiceDefinition` for this class."""
-    setattr(cls, "__nexus_service__", service_definition)
+    setattr(cls, _NEXUS_SERVICE_DEFINITION_ATTR_NAME, service_definition)
 
 
-def get_operation_definition(
+def get_operation(
     obj: Any,
 ) -> Optional[nexusrpc.Operation[Any, Any]]:
     """Return the :py:class:`nexusrpc.Operation` for the object, or None
 
     ``obj`` should be a decorated operation start method.
     """
-    return getattr(obj, "__nexus_operation__", None)
+    op = getattr(obj, _NEXUS_OPERATION_ATTR_NAME, None)
+    if op and not isinstance(op, nexusrpc.Operation):
+        raise ValueError(f"{op} is not a nexusrpc.Operation")
+    return op
 
 
-def set_operation_definition(
+def set_operation(
     obj: Any,
-    operation_definition: nexusrpc.Operation[Any, Any],
+    operation: nexusrpc.Operation[Any, Any],
 ) -> None:
     """Set the :py:class:`nexusrpc.Operation` for this object.
 
     ``obj`` should be an operation start method.
     """
-    setattr(obj, "__nexus_operation__", operation_definition)
+    if not isinstance(operation, nexusrpc.Operation):  # type: ignore
+        raise ValueError(f"{operation} is not a nexusrpc.Operation")  # type: ignore
+    setattr(obj, _NEXUS_OPERATION_ATTR_NAME, operation)
 
 
 def get_operation_factory(
@@ -72,15 +82,15 @@ def get_operation_factory(
 
     ``obj`` should be a decorated operation start method.
     """
-    op_defn = get_operation_definition(obj)
-    if op_defn:
+    op = get_operation(obj)
+    if op:
         factory = obj
     else:
-        if factory := getattr(obj, "__nexus_operation_factory__", None):
-            op_defn = get_operation_definition(factory)
-    if not isinstance(op_defn, nexusrpc.Operation):
+        if factory := getattr(obj, _NEXUS_OPERATION_FACTORY_ATTR_NAME, None):
+            op = get_operation(factory)
+    if not isinstance(op, nexusrpc.Operation):
         return None, None
-    return factory, op_defn
+    return factory, op
 
 
 def set_operation_factory(
@@ -91,7 +101,7 @@ def set_operation_factory(
 
     ``obj`` should be an operation start method.
     """
-    setattr(obj, "__nexus_operation_factory__", operation_factory)
+    setattr(obj, _NEXUS_OPERATION_FACTORY_ATTR_NAME, operation_factory)
 
 
 # Copied from https://github.com/modelcontextprotocol/python-sdk