Clarify get_operation_factory (#17)

Author: dandavison

src/nexusrpc/_util.py

@@ -51,8 +51,12 @@ def get_operation(
 ) -> Optional[nexusrpc.Operation[Any, Any]]:
     """Return the :py:class:`nexusrpc.Operation` for the object, or None
 
-    ``obj`` should be a decorated operation start method.
+    ``obj`` should be a decorated operation start method, or a method that takes
+    no arguments and returns an OperationHandler.
     """
+    if factory := getattr(obj, _NEXUS_OPERATION_FACTORY_ATTR_NAME, None):
+        # obj was a decorated operation start method
+        obj = factory
     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")
@@ -74,23 +78,15 @@ def set_operation(
 
 def get_operation_factory(
     obj: Any,
-) -> tuple[
-    Optional[Callable[[Any], OperationHandler[InputT, OutputT]]],
-    Optional[nexusrpc.Operation[InputT, OutputT]],
-]:
-    """Return the :py:class:`Operation` for the object along with the factory function.
-
-    ``obj`` should be a decorated operation start method.
-    """
-    op = get_operation(obj)
-    if op:
-        factory = obj
-    else:
-        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
+) -> Optional[Callable[[Any], OperationHandler[Any, Any]]]:
+    """Return the :py:class:`OperationHandler` factory function for the object."""
+    if factory := getattr(obj, _NEXUS_OPERATION_FACTORY_ATTR_NAME, None):
+        # obj was a decorated operation start method
+        return factory
+    if get_operation(obj):
+        # obj was the desired factory
+        return obj
+    return None
 
 
 def set_operation_factory(

src/nexusrpc/handler/_operation_handler.py

@@ -8,6 +8,7 @@
 from nexusrpc._common import InputT, OperationInfo, OutputT, ServiceHandlerT
 from nexusrpc._service import Operation, OperationDefinition, ServiceDefinition
 from nexusrpc._util import (
+    get_operation,
     get_operation_factory,
     is_async_callable,
     is_callable,
@@ -153,31 +154,28 @@ def collect_operation_handler_factories_by_method_name(
     )
     seen = set()
     for _, method in inspect.getmembers(user_service_cls, is_callable):
-        factory, op = get_operation_factory(method)  # type: ignore[var-annotated]
-        if factory and isinstance(op, Operation):
-            # This is a method decorated with one of the *operation_handler decorators
-            if op.name in seen:
-                raise RuntimeError(
-                    f"Operation '{op.name}' in service '{user_service_cls.__name__}' "
-                    f"is defined multiple times."
+        if factory := get_operation_factory(method):
+            if op := get_operation(factory):
+                if op.name in seen:
+                    raise RuntimeError(
+                        f"Operation '{op.name}' in service '{user_service_cls.__name__}' "
+                        f"is defined multiple times."
+                    )
+                if service and op.method_name not in service_method_names:
+                    _names = ", ".join(f"'{s}'" for s in sorted(service_method_names))
+                    msg = (
+                        f"Operation method name '{op.method_name}' in service handler {user_service_cls} "
+                        f"does not match an operation method name in the service definition. "
+                        f"Available method names in the service definition: "
+                    )
+                    msg += _names if _names else "[none]"
+                    raise TypeError(msg)
+
+                assert op.method_name, (
+                    f"Operation '{op}' method name should not be empty. Please report this as a bug."
                 )
-            if service and op.method_name not in service_method_names:
-                _names = ", ".join(f"'{s}'" for s in sorted(service_method_names))
-                msg = (
-                    f"Operation method name '{op.method_name}' in service handler {user_service_cls} "
-                    f"does not match an operation method name in the service definition. "
-                    f"Available method names in the service definition: "
-                )
-                msg += _names if _names else "[none]"
-                msg += "."
-                raise TypeError(msg)
-
-            # TODO(preview) op_defn.method name should be non-nullable
-            assert op.method_name, (
-                f"Operation '{op}' method name should not be None. This is an SDK bug."
-            )
-            factories[op.method_name] = factory
-            seen.add(op.name)
+                factories[op.method_name] = factory
+                seen.add(op.name)
     return factories
 
 
@@ -212,7 +210,7 @@ def validate_operation_handler_methods(
                 f"method name '{op_defn.method_name}'. But this operation is in service "
                 f"definition '{service_definition}'."
             )
-        _, op = get_operation_factory(factory)
+        op = get_operation(factory)
         if not isinstance(op, Operation):
             raise ValueError(
                 f"Method '{factory}' in class '{service_cls.__name__}' "
@@ -278,7 +276,7 @@ def service_definition_from_operation_handler_methods(
     """
     op_defns: dict[str, OperationDefinition[Any, Any]] = {}
     for name, method in user_methods.items():
-        _, op = get_operation_factory(method)
+        op = get_operation(method)
         if not isinstance(op, Operation):
             raise ValueError(
                 f"In service '{service_name}', could not locate operation definition for "

tests/handler/test_request_routing.py

@@ -5,7 +5,7 @@
 
 import nexusrpc
 from nexusrpc import LazyValue
-from nexusrpc._util import get_operation_factory, get_service_definition
+from nexusrpc._util import get_operation, get_service_definition
 from nexusrpc.handler import (
     Handler,
     StartOperationContext,
@@ -33,7 +33,7 @@ class UserServiceHandler:
         async def _op_impl(self, ctx: StartOperationContext, input: None) -> bool:
             assert (service_defn := get_service_definition(self.__class__))
             assert ctx.service == service_defn.name
-            _, op_handler_op_defn = get_operation_factory(self.op)
+            op_handler_op_defn = get_operation(self.op)
             assert op_handler_op_defn
             assert service_defn.operation_definitions.get(ctx.operation)
             return True

tests/handler/test_sync_operation_handler_decorator_creates_valid_operation_handler.py

@@ -35,7 +35,7 @@ async def my_async_def_op(self, ctx: StartOperationContext, input: int) -> int:
 
 def test_def_sync_handler():
     user_instance = MyServiceHandler()
-    op_handler_factory, _ = get_operation_factory(user_instance.my_def_op)
+    op_handler_factory = get_operation_factory(user_instance.my_def_op)
     assert op_handler_factory
     op_handler = op_handler_factory(user_instance)
     assert not is_async_callable(op_handler.start)
@@ -54,7 +54,7 @@ def test_def_sync_handler():
 @pytest.mark.asyncio
 async def test_async_def_sync_handler():
     user_instance = MyServiceHandler()
-    op_handler_factory, _ = get_operation_factory(user_instance.my_async_def_op)
+    op_handler_factory = get_operation_factory(user_instance.my_async_def_op)
     assert op_handler_factory
     op_handler = op_handler_factory(user_instance)
     assert is_async_callable(op_handler.start)