diff --git a/temporalio/contrib/streams/_outside.py b/temporalio/contrib/streams/_outside.py index 8f71aa9a8..c2454e7c8 100644 --- a/temporalio/contrib/streams/_outside.py +++ b/temporalio/contrib/streams/_outside.py @@ -102,6 +102,7 @@ def subscribe( from_offset: int | None = None, wait_new_records: bool = True, idle_timeout: timedelta | None = None, + type: Any = None, ) -> AsyncGenerator[Item, None]: """Read the Stream's items in order. @@ -114,6 +115,8 @@ def subscribe( from_offset: Start at this offset, inclusive. Defaults to the first Record. wait_new_records: Wait for new Records at the end of the Stream. idle_timeout: End with an ``Idle`` item when no new Record comes for this long. + type: Decode the values of Topics with no type as this type. A Topic from + :py:meth:`topic` with a type keeps its own. """ return self._read( topics=[topic.name for topic in topics] if topics else None, @@ -121,6 +124,7 @@ def subscribe( **self._types, **{topic.name: topic.type for topic in topics or []}, }, + default_type=type, from_offset=from_offset, wait_new_records=wait_new_records, idle_timeout=idle_timeout, @@ -131,6 +135,7 @@ async def _read( *, topics: list[str] | None, types: dict[str, Any], + default_type: Any, from_offset: int | None, wait_new_records: bool, idle_timeout: timedelta | None, @@ -141,6 +146,7 @@ async def _read( await self._stream_id(), topics=topics, types=types, + default_type=default_type, from_offset=from_offset, wait_new_records=wait_new_records, idle_timeout=idle_timeout, diff --git a/temporalio/contrib/streams/_reader.py b/temporalio/contrib/streams/_reader.py index 11f5e072e..05ad28df4 100644 --- a/temporalio/contrib/streams/_reader.py +++ b/temporalio/contrib/streams/_reader.py @@ -31,6 +31,7 @@ async def read( *, topics: Sequence[str] | None, types: Mapping[str, Any], + default_type: Any = None, from_offset: int | None, wait_new_records: bool, idle_timeout: timedelta | None, @@ -78,7 +79,7 @@ async def read( continue raise for record in answer.records or []: - item = await _item(converter, record, types) + item = await _item(converter, record, types, default_type) if item is None: continue yield item @@ -98,11 +99,12 @@ async def _item( converter: temporalio.converter.DataConverter, record: models.StreamRecord, types: Mapping[str, Any], + default_type: Any, ) -> Item | None: offset = record.offset or 0 if record.kind == StreamRecordKind.STREAM_RECORD_KIND_DATA: topic = record.topic or "" - type_hint = types.get(topic) + type_hint = types.get(topic) or default_type data = await _convert.decode(converter, record.body, type_hint) return Record(topic=topic, data=data, offset=offset, producer=record.producer) if record.kind == StreamRecordKind.STREAM_RECORD_KIND_SUPERSEDED: diff --git a/temporalio/contrib/streams/_stream.py b/temporalio/contrib/streams/_stream.py index 3c153733b..6d426d9b5 100644 --- a/temporalio/contrib/streams/_stream.py +++ b/temporalio/contrib/streams/_stream.py @@ -42,6 +42,7 @@ def subscribe( from_offset: int | None = None, wait_new_records: bool = True, idle_timeout: timedelta | None = None, + type: Any = None, ) -> AsyncGenerator[Item, None]: ... diff --git a/tests/contrib/streams/test_facade.py b/tests/contrib/streams/test_facade.py index 54c2c2f9d..f30b8192a 100644 --- a/tests/contrib/streams/test_facade.py +++ b/tests/contrib/streams/test_facade.py @@ -306,3 +306,18 @@ async def test_create_wait_ends_with_the_run(client: Client): start = asyncio.get_running_loop().time() assert [item async for item in stream.subscribe()] == [] assert asyncio.get_running_loop().time() - start < 30 + + +async def test_read_default_type_for_untyped_topics(client: Client): + client, _ = fake_client(client) + stream = await streams.create_stream(client, "default-type") + stream.topic("turns").publish(Turn("model")) + stream.topic("tokens").publish("hi") + await stream.flush() + reader = streams.for_ref(client, stream.ref) + reader.topic("tokens", type=RawValue) + items = [item async for item in reader.subscribe(type=Turn, wait_new_records=False)] + first, second = items + assert isinstance(first, streams.Record) and first.data == Turn("model") + # A typed Topic keeps its own type. + assert isinstance(second, streams.Record) and isinstance(second.data, RawValue)