From a6ff7bebd2d413fed94aef7d39bbb02776218d0a Mon Sep 17 00:00:00 2001 From: Mohammad Dashti Date: Thu, 22 Oct 2026 14:09:06 -0700 Subject: [PATCH] Named a Stream by a renamed definition, not by a name option. --- temporalio/contrib/streams/__init__.py | 2 + temporalio/contrib/streams/_defn.py | 37 ++++++++++++++-- temporalio/contrib/streams/_stream_client.py | 6 +-- temporalio/contrib/streams/workflow.py | 36 ++++++---------- .../streams/examples/resume_from_offset.py | 11 +++-- .../streams/examples/turn_scoped_stream.py | 5 +-- .../contrib/streams/test_activity_handles.py | 6 ++- tests/contrib/streams/test_redis.py | 14 ++++--- tests/contrib/streams/test_stream_client.py | 4 +- tests/contrib/streams/test_workflow_module.py | 42 ++++++++++++++++++- 10 files changed, 114 insertions(+), 49 deletions(-) diff --git a/temporalio/contrib/streams/__init__.py b/temporalio/contrib/streams/__init__.py index 289bf22da..f60569b5b 100644 --- a/temporalio/contrib/streams/__init__.py +++ b/temporalio/contrib/streams/__init__.py @@ -31,6 +31,7 @@ TopicDefinition, defn, topic, + with_name, ) from temporalio.contrib.streams._errors import ( ReadOnlyContextError, @@ -85,5 +86,6 @@ "WorkflowTopic", "defn", "topic", + "with_name", "workflow", ] diff --git a/temporalio/contrib/streams/_defn.py b/temporalio/contrib/streams/_defn.py index 86a99e765..2708a5fe1 100644 --- a/temporalio/contrib/streams/_defn.py +++ b/temporalio/contrib/streams/_defn.py @@ -3,8 +3,8 @@ from __future__ import annotations from collections.abc import Callable, Mapping -from dataclasses import dataclass -from typing import Any, Generic, TypeVar, overload +from dataclasses import dataclass, replace +from typing import Any, Generic, TypeVar, cast, overload T = TypeVar("T") ClassT = TypeVar("ClassT", bound=type) @@ -82,7 +82,7 @@ class StreamDefinition: """ name: str - """Default name of a Stream of this definition.""" + """Name of a Stream of this definition.""" topics: Mapping[str, TopicDefinition[Any]] """The Topics by name.""" @@ -105,7 +105,8 @@ def defn( Args: cls: The class. - name: The default name of a Stream of this definition. Defaults to the class name. + name: The name of a Stream of this definition. Defaults to the class name. Use + :py:func:`with_name` for another Stream of the same Topics. .. warning:: This API is experimental. @@ -138,6 +139,34 @@ def decorate(cls: ClassT) -> ClassT: return decorate +def with_name(definition: ClassT, name: str) -> ClassT: + """The definition with the same Topics and another Stream name, such as one per turn. + + It is a subclass of ``definition``, which stays unchanged. Its Topics are the same + attributes, so ``Events.tokens`` is a Topic of ``with_name(Events, "turn-3")`` too. + + Args: + definition: The definition class. + name: The Stream name. + + .. warning:: + This API is experimental. + """ + renamed = replace(definition_of(definition), name=name) + return cast( + ClassT, + type( + definition.__name__, + (definition,), + { + _DEFINITION_ATTR: renamed, + "__module__": definition.__module__, + "__qualname__": definition.__qualname__, + }, + ), + ) + + def definition_of(cls: type) -> StreamDefinition: """The definition that :py:func:`defn` gave a class. diff --git a/temporalio/contrib/streams/_stream_client.py b/temporalio/contrib/streams/_stream_client.py index ac878a569..bc3564e4f 100644 --- a/temporalio/contrib/streams/_stream_client.py +++ b/temporalio/contrib/streams/_stream_client.py @@ -370,7 +370,6 @@ def workflow_stream( handle: temporalio.client.WorkflowHandle[Any, Any], definition: type[E], *, - name: str | None = None, backend: str | tuple[str, ...] | None = None, ) -> ClientStream[E]: """A Stream that a Workflow created, by the Workflow handle and the definition. @@ -380,12 +379,11 @@ def workflow_stream( Args: handle: A handle of the Workflow, best the one that ``start_workflow`` returned. - definition: The Stream definition class. - name: The Stream name that the Workflow gave. Defaults to the definition's name. + definition: The Stream definition class, renamed as the Workflow's create. backend: The name of the backend, or a tuple of names to look in, as for :py:meth:`get_handle`. Defaults to the client's default backend. """ - stream_name = name or definition_of(definition).name + stream_name = definition_of(definition).name backend_name, stream_backend, probe = self._backends(backend) first_run = handle.first_execution_run_id stream = ClientStream( diff --git a/temporalio/contrib/streams/workflow.py b/temporalio/contrib/streams/workflow.py index 87803aa0e..63ac5c233 100644 --- a/temporalio/contrib/streams/workflow.py +++ b/temporalio/contrib/streams/workflow.py @@ -346,14 +346,13 @@ async def __aexit__( await self.close(failure=exc) -def _own_handle( - definition: type[E], name: str | None, backend: str | None -) -> WorkflowStream[E]: +def _own_handle(definition: type[E], backend: str | None) -> WorkflowStream[E]: info = temporalio.workflow.info() - stream_name = name or definition_of(definition).name return WorkflowStream( id=_ids.scoped_stream_id( - info.workflow_id, info.first_execution_run_id, stream_name + info.workflow_id, + info.first_execution_run_id, + definition_of(definition).name, ), workflow_id=info.workflow_id, first_execution_run_id=info.first_execution_run_id, @@ -365,7 +364,6 @@ def _own_handle( async def create_stream( definition: type[E], *, - name: str | None = None, retention: timedelta | None = None, auto_close: timedelta | None = None, close_with_run: bool = True, @@ -377,9 +375,8 @@ async def create_stream( handle with no call. The create is one Local Activity, and its error raises here. Args: - definition: The Stream definition class. - name: The Stream name. Defaults to the definition's name. Give one Stream per turn, - for example, its own name. + definition: The Stream definition class. Its name is the Stream name. For one Stream + per turn, for example, pass ``streams.with_name(Events, f"turn-{n}")``. retention: How long the Records stay after the close. Defaults to the backend's. auto_close: Close the Stream with no result when no Record was appended for this long, so a Stream that nobody closes does not stay open. Defaults to never. @@ -398,9 +395,7 @@ async def create_stream( ValueError: ``close_with_run=False`` for a Stream that already closes with the run chain. """ - stream, task = _create( - definition, name, retention, auto_close, close_with_run, backend - ) + stream, task = _create(definition, retention, auto_close, close_with_run, backend) await task return stream @@ -408,7 +403,6 @@ async def create_stream( def create_stream_nowait( definition: type[E], *, - name: str | None = None, retention: timedelta | None = None, auto_close: timedelta | None = None, close_with_run: bool = True, @@ -420,22 +414,19 @@ def create_stream_nowait( is for code that can't await ``put``. It takes the options of :py:func:`create_stream`. A backend error fails the Workflow Task. """ - stream, task = _create( - definition, name, retention, auto_close, close_with_run, backend - ) + stream, task = _create(definition, retention, auto_close, close_with_run, backend) _write_nowait(task) return stream def _create( definition: type[E], - name: str | None, retention: timedelta | None, auto_close: timedelta | None, close_with_run: bool, backend: str | None, ) -> tuple[WorkflowStream[E], asyncio.Task[int]]: - stream = _own_handle(definition, name, _check_backend(backend, "create")) + stream = _own_handle(definition, _check_backend(backend, "create")) run_stream = stream._check("create", backend_checked=True) if run_stream.close_with_run and not close_with_run: raise ValueError( @@ -462,20 +453,19 @@ def _create( def stream_handle( - definition: type[E], *, name: str | None = None, backend: str | None = None + definition: type[E], *, backend: str | None = None ) -> WorkflowStream[E]: - """A Stream that this run chain created, by its definition and name. It makes no call. + """A Stream that this run chain created, by its definition. It makes no call. A later run of the chain, such as one after continue-as-new, uses it to write the Stream again. Args: - definition: The Stream definition class. - name: The Stream name. Defaults to the definition's name. + definition: The Stream definition class, renamed as for the create. backend: The backend of the Stream. Defaults to the one that the run chain created it on, and else to the Worker's. """ - stream = _own_handle(definition, name, backend) + stream = _own_handle(definition, backend) if backend is None: known = _state().streams.get(stream.id) stream._backend = (known.backend if known else None) or _backend_name() or None diff --git a/tests/contrib/streams/examples/resume_from_offset.py b/tests/contrib/streams/examples/resume_from_offset.py index 955d7d7fd..8bc7a722b 100644 --- a/tests/contrib/streams/examples/resume_from_offset.py +++ b/tests/contrib/streams/examples/resume_from_offset.py @@ -1,6 +1,6 @@ """Brook API proposal, "A client resumes a scoped Stream from an offset". -The Workflow writes one Stream per turn with a ``name=`` override. The Reader saves the offset +The Workflow writes one Stream per turn, a renamed definition from ``streams.with_name``. The Reader saves the offset after each Record itself, stops, and later resumes from it with ``from_offset``. A view that wants only new Records starts at ``streams.END``. """ @@ -22,6 +22,9 @@ class Events: tokens = streams.topic(str) +Turn3 = streams.with_name(Events, "turn-3") + + @workflow.defn class Turns: def __init__(self) -> None: @@ -30,7 +33,7 @@ def __init__(self) -> None: @workflow.run async def run(self, first: list[str]) -> None: - turn = await stream_workflow.create_stream(Events, name="turn-3") + turn = await stream_workflow.create_stream(Turn3) tokens = turn.topic(Events.tokens) await tokens.publish(*first) await workflow.wait_condition(lambda: bool(self.more) or self.done) @@ -57,7 +60,7 @@ def save(self, key: str, offset: int) -> None: async def read_some( client: Client, handle: WorkflowHandle[Turns, None], store: OffsetStore, count: int ) -> list[str]: - turn3 = streams.StreamClient(client).workflow_stream(handle, Events, name="turn-3") + turn3 = streams.StreamClient(client).workflow_stream(handle, Turn3) tokens = turn3.topic(Events.tokens) seen = [] async for item in tokens.subscribe( @@ -77,7 +80,7 @@ async def read_some( async def read_new(client: Client, handle: WorkflowHandle[Turns, None]) -> list[str]: # A view that shows only what comes next starts after the newest Record. - turn3 = streams.StreamClient(client).workflow_stream(handle, Events, name="turn-3") + turn3 = streams.StreamClient(client).workflow_stream(handle, Turn3) return [ item.data async for item in turn3.topic(Events.tokens).subscribe( diff --git a/tests/contrib/streams/examples/turn_scoped_stream.py b/tests/contrib/streams/examples/turn_scoped_stream.py index e1640e46c..0b2e5a0c9 100644 --- a/tests/contrib/streams/examples/turn_scoped_stream.py +++ b/tests/contrib/streams/examples/turn_scoped_stream.py @@ -1,6 +1,6 @@ """Brook API proposal, "A Workflow scopes a Stream to one turn and closes it". -Each turn gets its own Stream, named with ``name=``, which ``async with`` closes at the end of +Each turn gets its own Stream, named with ``streams.with_name``, which ``async with`` closes at the end of the turn. An index Stream lists the turns, and continue-as-new passes it to the next run. """ @@ -64,8 +64,7 @@ async def run(self, input: ChatTurnsInput) -> None: batch = input.prompts[:2] for i, prompt in enumerate(batch): async with await stream_workflow.create_stream( - Events, - name=f"turn-{input.next + i}", + streams.with_name(Events, f"turn-{input.next + i}"), retention=timedelta(days=1), close_with_run=False, ) as turn: diff --git a/tests/contrib/streams/test_activity_handles.py b/tests/contrib/streams/test_activity_handles.py index fd34e21f5..0b74f6ebb 100644 --- a/tests/contrib/streams/test_activity_handles.py +++ b/tests/contrib/streams/test_activity_handles.py @@ -47,7 +47,9 @@ class SyncInput: async def read( client: Client, handle: WorkflowHandle[Any, Any], name: str = "events" ) -> list[object]: - stream = streams.StreamClient(client).workflow_stream(handle, Events, name=name) + stream = streams.StreamClient(client).workflow_stream( + handle, streams.with_name(Events, name) + ) return [item async for item in stream.subscribe(create_wait=timedelta(seconds=20))] @@ -362,7 +364,7 @@ class Leftovers: @workflow.run async def run(self, key: str) -> list[Any]: stream = await stream_workflow.create_stream(Events) - other = await stream_workflow.create_stream(Events, name="other") + other = await stream_workflow.create_stream(streams.with_name(Events, "other")) await workflow.execute_activity( never_flushes, CallModelInput("a b", streams.ActivityStream.from_handle(stream)), diff --git a/tests/contrib/streams/test_redis.py b/tests/contrib/streams/test_redis.py index a20a416d6..150832009 100644 --- a/tests/contrib/streams/test_redis.py +++ b/tests/contrib/streams/test_redis.py @@ -59,7 +59,9 @@ class CallModelInput: def workflow_stream( client: Client, handle: WorkflowHandle[Any, Any], name: str | None = None ) -> streams.ClientStream[Events]: - return streams.StreamClient(client).workflow_stream(handle, Events, name=name) + return streams.StreamClient(client).workflow_stream( + handle, streams.with_name(Events, name) if name else Events + ) async def read_all(stream: streams.ClientStream[Any]) -> list[Any]: @@ -220,7 +222,7 @@ class CancelledBlock: @workflow.run async def run(self) -> None: async with await stream_workflow.create_stream( - Events, name="turn", close_with_run=False + streams.with_name(Events, "turn"), close_with_run=False ) as turn: await turn.topic(Events.tokens).publish("started") await workflow.wait_condition(lambda: False) @@ -254,7 +256,7 @@ async def run(self) -> str: stream = await stream_workflow.create_stream(Plain) await stream.topic(Plain.t).publish("plain") async with await stream_workflow.create_stream( - Plain, name="scoped", close_with_run=False + streams.with_name(Plain, "scoped"), close_with_run=False ) as scoped: await scoped.topic(Plain.t).publish("scoped") return "done" @@ -276,7 +278,9 @@ async def test_temporal_only_passthrough_keeps_records(stream_client: Client): await handle.result() plain = streams.StreamClient(stream_client) items = await read_all(plain.workflow_stream(handle, Plain)) - scoped = await read_all(plain.workflow_stream(handle, Plain, name="scoped")) + scoped = await read_all( + plain.workflow_stream(handle, streams.with_name(Plain, "scoped")) + ) assert [i.data for i in records(items)] == ["plain"] assert isinstance(items[-1], streams.Closed) and items[-1].result == "done" assert [i.data for i in records(scoped)] == ["scoped"] @@ -517,7 +521,7 @@ async def run(self) -> str: turns = (await stream_workflow.create_stream(Events)).topic(Events.tokens) await turns.publish("before restart") async with await stream_workflow.create_stream( - Events, name="scoped", close_with_run=False + streams.with_name(Events, "scoped"), close_with_run=False ) as scoped: await scoped.topic(Events.tokens).publish("scoped") await workflow.wait_condition(lambda: self.go) diff --git a/tests/contrib/streams/test_stream_client.py b/tests/contrib/streams/test_stream_client.py index 9c499497e..ee5aeb3d9 100644 --- a/tests/contrib/streams/test_stream_client.py +++ b/tests/contrib/streams/test_stream_client.py @@ -94,7 +94,7 @@ class Description: setattr(handle, "describe", describe) stream_client = streams.StreamClient(client) await stream_client.create(Events, id="agent-1:first:n") - stream = stream_client.workflow_stream(handle, Events, name="n") + stream = stream_client.workflow_stream(handle, streams.with_name(Events, "n")) with pytest.raises(RuntimeError, match="not known yet"): stream.id assert await stream.get_offset() == 0 @@ -200,7 +200,7 @@ async def test_unresolved_workflow_stream_compares_and_prints(client: Client): handle = client.get_workflow_handle("agent-2") first = stream_client.workflow_stream(handle, Events) same = stream_client.workflow_stream(handle, Events) - other = stream_client.workflow_stream(handle, Events, name="other") + other = stream_client.workflow_stream(handle, streams.with_name(Events, "other")) key = hash(first) assert first == same and first != other describes = [] diff --git a/tests/contrib/streams/test_workflow_module.py b/tests/contrib/streams/test_workflow_module.py index 603e04349..6633494a5 100644 --- a/tests/contrib/streams/test_workflow_module.py +++ b/tests/contrib/streams/test_workflow_module.py @@ -15,6 +15,7 @@ ) from temporalio.contrib import streams from temporalio.contrib.streams import workflow as stream_workflow +from temporalio.contrib.streams._defn import definition_of from temporalio.contrib.streams.system import client as service_client from temporalio.exceptions import ApplicationError from temporalio.worker import Replayer @@ -36,7 +37,9 @@ class Events: async def read( client: Client, handle: WorkflowHandle[Any, Any], name: str = "events" ) -> list[object]: - stream = streams.StreamClient(client).workflow_stream(handle, Events, name=name) + stream = streams.StreamClient(client).workflow_stream( + handle, streams.with_name(Events, name) + ) return [item async for item in stream.subscribe(create_wait=timedelta(seconds=20))] @@ -74,7 +77,7 @@ async def run(self, fail_at: int) -> None: for n in range(2): try: async with await stream_workflow.create_stream( - Events, name=f"turn-{n}", close_with_run=False + streams.with_name(Events, f"turn-{n}"), close_with_run=False ) as turn: await turn.topic(Events.tokens).publish(f"token-{n}") if n == fail_at: @@ -113,6 +116,41 @@ async def run(self, stream: stream_workflow.WorkflowStream[Events] | None) -> st return "done" +@workflow.defn +class RenamedAcrossRuns: + @workflow.run + async def run(self, first: bool) -> list[str]: + renamed = streams.with_name(Events, "turn-1") + if first: + turn = await stream_workflow.create_stream(renamed) + await turn.topic(Events.tokens).publish("run 1") + workflow.continue_as_new(False) + turn = stream_workflow.stream_handle(renamed) + await turn.topic(Events.tokens).publish("run 2") + return [stream_workflow.stream_handle(Events).id, turn.id] + + +async def test_renamed_definition_is_another_stream(stream_client: Client): + async with new_worker(stream_client, RenamedAcrossRuns) as worker: + handle = await stream_client.start_workflow( + RenamedAcrossRuns.run, + True, + id=f"wf-{uuid.uuid4()}", + task_queue=worker.task_queue, + ) + plain_id, turn_id = await handle.result() + turn = await read(stream_client, handle, "turn-1") + assert plain_id.endswith(":events") and turn_id.endswith(":turn-1") + assert data(turn) == ["run 1", "run 2"] + + +def test_with_name_leaves_the_definition_unchanged(): + renamed = streams.with_name(Events, "turn-1") + assert definition_of(Events).name == "events" + assert definition_of(renamed).name == "turn-1" + assert definition_of(renamed).topics == definition_of(Events).topics + + async def test_continue_as_new_passes_the_stream(stream_client: Client): async with new_worker(stream_client, Continuing) as worker: handle = await stream_client.start_workflow(