Make Temporal context classes non-private

Author: dandavison

temporalio/nexus/__init__.py

@@ -1,13 +1,13 @@
 from ._decorators import workflow_run_operation as workflow_run_operation
 from ._operation_context import Info as Info
 from ._operation_context import (
-    WorkflowRunOperationContext as WorkflowRunOperationContext,
+    TemporalStartOperationContext as TemporalStartOperationContext,
 )
 from ._operation_context import (
     _TemporalCancelOperationContext as _TemporalCancelOperationContext,
 )
 from ._operation_context import (
-    _TemporalStartOperationContext as _TemporalStartOperationContext,
+    WorkflowRunOperationContext as WorkflowRunOperationContext,
 )
 from ._operation_context import client as client
 from ._operation_context import info as info

temporalio/nexus/_decorators.py

@@ -17,8 +17,8 @@
 )
 
 from temporalio.nexus._operation_context import (
+    TemporalStartOperationContext,
     WorkflowRunOperationContext,
-    _TemporalStartOperationContext,
 )
 from temporalio.nexus._operation_handlers import (
     WorkflowRunOperationHandler,
@@ -115,7 +115,7 @@ def operation_handler_factory(
             async def _start(
                 ctx: StartOperationContext, input: InputT
             ) -> WorkflowHandle[OutputT]:
-                tctx = _TemporalStartOperationContext.get()
+                tctx = TemporalStartOperationContext.get()
                 return await start(self, WorkflowRunOperationContext(tctx), input)
 
             _start.__doc__ = start.__doc__

temporalio/nexus/_operation_context.py

@@ -35,7 +35,7 @@
 # CancelOperationContext and passes it as the first parameter to the nexusrpc operation
 # handler. In addition, it sets one of the following context vars.
 
-_temporal_start_operation_context: ContextVar[_TemporalStartOperationContext] = (
+_temporal_start_operation_context: ContextVar[TemporalStartOperationContext] = (
     ContextVar("temporal-start-operation-context")
 )
 
@@ -70,7 +70,7 @@ def client() -> temporalio.client.Client:
 
 
 def _temporal_context() -> (
-    Union[_TemporalStartOperationContext, _TemporalCancelOperationContext]
+    Union[TemporalStartOperationContext, _TemporalCancelOperationContext]
 ):
     ctx = _try_temporal_context()
     if ctx is None:
@@ -79,7 +79,7 @@ def _temporal_context() -> (
 
 
 def _try_temporal_context() -> (
-    Optional[Union[_TemporalStartOperationContext, _TemporalCancelOperationContext]]
+    Optional[Union[TemporalStartOperationContext, _TemporalCancelOperationContext]]
 ):
     start_ctx = _temporal_start_operation_context.get(None)
     cancel_ctx = _temporal_cancel_operation_context.get(None)
@@ -89,7 +89,7 @@ def _try_temporal_context() -> (
 
 
 @dataclass
-class _TemporalStartOperationContext:
+class TemporalStartOperationContext:
     """
     Context for a Nexus start operation being handled by a Temporal Nexus Worker.
     """
@@ -104,7 +104,7 @@ class _TemporalStartOperationContext:
     """The Temporal client in use by the worker handling this Nexus operation."""
 
     @classmethod
-    def get(cls) -> _TemporalStartOperationContext:
+    def get(cls) -> TemporalStartOperationContext:
         ctx = _temporal_start_operation_context.get(None)
         if ctx is None:
             raise RuntimeError("Not in Nexus operation context.")
@@ -170,15 +170,15 @@ def add_outbound_links(
 
 @dataclass
 class WorkflowRunOperationContext:
-    temporal_context: _TemporalStartOperationContext
+    temporal_context: TemporalStartOperationContext
 
     @property
     def nexus_context(self) -> StartOperationContext:
         return self.temporal_context.nexus_context
 
     @classmethod
     def get(cls) -> WorkflowRunOperationContext:
-        return cls(_TemporalStartOperationContext.get())
+        return cls(TemporalStartOperationContext.get())
 
     # Overload for single-param workflow
     # TODO(nexus-prerelease): bring over other overloads

temporalio/worker/_nexus.py

@@ -34,8 +34,8 @@
 from temporalio.exceptions import ApplicationError
 from temporalio.nexus import (
     Info,
+    TemporalStartOperationContext,
     _TemporalCancelOperationContext,
-    _TemporalStartOperationContext,
     logger,
 )
 from temporalio.service import RPCError, RPCStatusCode
@@ -269,7 +269,7 @@ async def _start_operation(
             ],
             callback_headers=dict(start_request.callback_header),
         )
-        _TemporalStartOperationContext(
+        TemporalStartOperationContext(
             nexus_context=ctx,
             client=self._client,
             info=lambda: Info(task_queue=self._task_queue),