Mark TemporalStartOperationContext as 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 (
-    TemporalStartOperationContext as TemporalStartOperationContext,
+    WorkflowRunOperationContext as WorkflowRunOperationContext,
 )
 from ._operation_context import (
     _TemporalCancelOperationContext as _TemporalCancelOperationContext,
 )
 from ._operation_context import (
-    WorkflowRunOperationContext as WorkflowRunOperationContext,
+    _TemporalStartOperationContext as _TemporalStartOperationContext,
 )
 from ._operation_context import client as client
 from ._operation_context import info as info

temporalio/nexus/_operation_context.py

@@ -36,7 +36,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")
 )
 
@@ -71,7 +71,7 @@ def client() -> temporalio.client.Client:
 
 
 def _temporal_context() -> (
-    Union[TemporalStartOperationContext, _TemporalCancelOperationContext]
+    Union[_TemporalStartOperationContext, _TemporalCancelOperationContext]
 ):
     ctx = _try_temporal_context()
     if ctx is None:
@@ -80,7 +80,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)
@@ -90,7 +90,7 @@ def _try_temporal_context() -> (
 
 
 @dataclass
-class TemporalStartOperationContext:
+class _TemporalStartOperationContext:
     """
     Context for a Nexus start operation being handled by a Temporal Nexus Worker.
     """
@@ -105,7 +105,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.")
@@ -171,10 +171,10 @@ def add_outbound_links(
 
 @dataclass(frozen=True)
 class WorkflowRunOperationContext(StartOperationContext):
-    _temporal_context: Optional[TemporalStartOperationContext] = None
+    _temporal_context: Optional[_TemporalStartOperationContext] = None
 
     @property
-    def temporal_context(self) -> TemporalStartOperationContext:
+    def temporal_context(self) -> _TemporalStartOperationContext:
         if not self._temporal_context:
             raise RuntimeError("Temporal context not set")
         return self._temporal_context
@@ -188,7 +188,7 @@ def from_start_operation_context(
         cls, ctx: StartOperationContext
     ) -> WorkflowRunOperationContext:
         return cls(
-            _temporal_context=TemporalStartOperationContext.get(),
+            _temporal_context=_TemporalStartOperationContext.get(),
             **{f.name: getattr(ctx, f.name) for f in dataclasses.fields(ctx)},
         )
 

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),