Factor out helper to share start-workflow code

Author: dandavison

temporalio/client.py

@@ -494,22 +494,12 @@ async def start_workflow(
                 already been started.
             RPCError: Workflow could not be started for some other reason.
         """
-        # Use definition if callable
-        name: str
-        if isinstance(workflow, str):
-            name = workflow
-        elif callable(workflow):
-            defn = temporalio.workflow._Definition.must_from_run_fn(workflow)
-            if not defn.name:
-                raise ValueError("Cannot invoke dynamic workflow explicitly")
-            name = defn.name
-            if result_type is None:
-                result_type = defn.ret_type
-        else:
-            raise TypeError("Workflow must be a string or callable")
         temporalio.common._warn_on_deprecated_search_attributes(
             search_attributes, stack_level=stack_level
         )
+        name, result_type_from_run_fn = (
+            temporalio.workflow._Definition.get_name_and_result_type(workflow)
+        )
 
         return await self._impl.start_workflow(
             StartWorkflowInput(
@@ -530,7 +520,7 @@ async def start_workflow(
                 headers={},
                 start_signal=start_signal,
                 start_signal_args=start_signal_args,
-                ret_type=result_type,
+                ret_type=result_type or result_type_from_run_fn,
                 rpc_metadata=rpc_metadata,
                 rpc_timeout=rpc_timeout,
                 request_eager_start=request_eager_start,
@@ -2321,7 +2311,9 @@ def get_update_handle_for(
 class WithStartWorkflowOperation(Generic[SelfType, ReturnType]):
     """
     Defines a start-workflow operation used by update-with-start requests.
-    Usually created by :py:meth:`Client.create_start_workflow_operation`.
+
+    Update-With-Start allows you to send an update to a workflow, while starting the
+    workflow if necessary.
     """
 
     # Overload for no-param workflow, with_start
@@ -2451,7 +2443,6 @@ def __init__(
         request_eager_start: bool = False,
     ) -> None: ...
 
-    # TODO: reduce duplication with start_workflow
     def __init__(
         self,
         workflow: Union[str, Callable[..., Awaitable[Any]]],
@@ -2484,30 +2475,17 @@ def __init__(
         stack_level: int = 2,
     ) -> None:
         """
-        Create a handle for issuing an Update-With-Start request.
+        Create a WithStartWorkflowOperation.
 
-        Update-With-Start allows you to send an update to a workflow, while starting the workflow
-        if necessary.
+        See :py:meth:`temporalio.client.Client.start_workflow` for documentation of the
+        arguments.
         """
-
-        # TODO: avoid duplicating this?
-
-        # Use definition if callable
-        name: str
-        if isinstance(workflow, str):
-            name = workflow
-        elif callable(workflow):
-            defn = temporalio.workflow._Definition.must_from_run_fn(workflow)
-            if not defn.name:
-                raise ValueError("Cannot invoke dynamic workflow explicitly")
-            name = defn.name
-            if result_type is None:
-                result_type = defn.ret_type
-        else:
-            raise TypeError("Workflow must be a string or callable")
         temporalio.common._warn_on_deprecated_search_attributes(
             search_attributes, stack_level=stack_level
         )
+        name, result_type_from_run_fn = (
+            temporalio.workflow._Definition.get_name_and_result_type(workflow)
+        )
 
         self._id = id
         self._workflow_handle: Future[WorkflowHandle[SelfType, ReturnType]] = Future()
@@ -2529,7 +2507,7 @@ def __init__(
             headers={},
             start_signal=start_signal,
             start_signal_args=start_signal_args,
-            ret_type=result_type,
+            ret_type=result_type or result_type_from_run_fn,
             rpc_metadata=rpc_metadata,
             rpc_timeout=rpc_timeout,
             request_eager_start=request_eager_start,

temporalio/workflow.py

@@ -1335,6 +1335,20 @@ def must_from_run_fn(fn: Callable[..., Awaitable[Any]]) -> _Definition:
             f"Function {fn_name} missing attributes, was it decorated with @workflow.run and was its class decorated with @workflow.defn?"
         )
 
+    @classmethod
+    def get_name_and_result_type(
+        cls, name_or_run_fn: Union[str, Callable[..., Awaitable[Any]]]
+    ) -> Tuple[str, Optional[Type]]:
+        if isinstance(name_or_run_fn, str):
+            return name_or_run_fn, None
+        elif callable(name_or_run_fn):
+            defn = cls.must_from_run_fn(name_or_run_fn)
+            if not defn.name:
+                raise ValueError("Cannot invoke dynamic workflow explicitly")
+            return defn.name, defn.ret_type
+        else:
+            raise TypeError("Workflow must be a string or callable")
+
     @staticmethod
     def _apply_to_class(
         cls: Type,