diff --git a/temporalio/contrib/streams/__init__.py b/temporalio/contrib/streams/__init__.py index acf5cdf947..9801b8d202 100644 --- a/temporalio/contrib/streams/__init__.py +++ b/temporalio/contrib/streams/__init__.py @@ -17,6 +17,12 @@ for_handle, for_ref, ) +from temporalio.contrib.streams._defn import ( + StreamDefinition, + TopicDefinition, + defn, + topic, +) from temporalio.contrib.streams._errors import ( ReadOnlyContextError, StreamBufferFullError, @@ -30,6 +36,7 @@ StreamTransportMissingError, StreamUnsupportedError, ) +from temporalio.contrib.streams._handle import StreamHandle from temporalio.contrib.streams._items import Closed, Record, Superseded from temporalio.contrib.streams._outside import OutsideStream from temporalio.contrib.streams._stream import StreamRef, Topic @@ -47,10 +54,12 @@ "Record", "StreamBackend", "StreamBufferFullError", + "StreamDefinition", "StreamClosedError", "StreamError", "StreamExpiredError", "StreamFencedError", + "StreamHandle", "StreamNotFoundError", "StreamOffsetError", "StreamRef", @@ -59,11 +68,14 @@ "StreamUnsupportedError", "Superseded", "Topic", + "TopicDefinition", "WorkflowStream", "activity_stream", "create_stream", "create_workflow_stream", + "defn", "for_handle", "for_ref", + "topic", "workflow_stream", ] diff --git a/temporalio/contrib/streams/_defn.py b/temporalio/contrib/streams/_defn.py new file mode 100644 index 0000000000..86a99e7653 --- /dev/null +++ b/temporalio/contrib/streams/_defn.py @@ -0,0 +1,150 @@ +"""Stream definitions: the Stream name and the type of each Topic.""" + +from __future__ import annotations + +from collections.abc import Callable, Mapping +from dataclasses import dataclass +from typing import Any, Generic, TypeVar, overload + +T = TypeVar("T") +ClassT = TypeVar("ClassT", bound=type) + +_DEFINITION_ATTR = "__temporal_stream_definition" + + +class TopicDefinition(Generic[T]): + """One Topic of a Stream definition. Use :py:func:`topic` in the definition class. + + .. warning:: + This API is experimental. + """ + + def __init__(self, type: Any, name: str | None = None) -> None: + """Create a Topic definition. Use :py:func:`topic` instead.""" + self._type = type + self._name = name + self._owner: type | None = None + + def __set_name__(self, owner: type, name: str) -> None: + """Take the attribute name as the Topic name, unless the Topic has a name.""" + if self._name is None: + self._name = name + self._owner = owner + + @property + def name(self) -> str: + """Name of the Topic.""" + if self._name is None: + raise TypeError( + "A Topic definition must be an attribute of a definition class" + ) + return self._name + + @property + def type(self) -> Any: + """Type that a read decodes the Topic's values with.""" + return self._type + + def __repr__(self) -> str: + """Show the Topic name and type.""" + return f"TopicDefinition(name={self._name!r}, type={self._type!r})" + + +@overload +def topic(type: type[T], *, name: str | None = None) -> TopicDefinition[T]: ... + + +@overload +def topic(type: Any, *, name: str | None = None) -> TopicDefinition[Any]: ... + + +def topic(type: Any, *, name: str | None = None) -> TopicDefinition[Any]: + """Declare a Topic in a Stream definition class. + + ``RawValue`` keeps each Record's Payload. + + Args: + type: The type of the Topic's values. + name: The Topic name. Defaults to the attribute name. + + .. warning:: + This API is experimental. + """ + return TopicDefinition(type, name) + + +@dataclass(frozen=True) +class StreamDefinition: + """What :py:func:`defn` takes from a definition class. + + .. warning:: + This API is experimental. + """ + + name: str + """Default name of a Stream of this definition.""" + + topics: Mapping[str, TopicDefinition[Any]] + """The Topics by name.""" + + +@overload +def defn(cls: ClassT) -> ClassT: ... + + +@overload +def defn(*, name: str | None = None) -> Callable[[ClassT], ClassT]: ... + + +def defn( + cls: ClassT | None = None, *, name: str | None = None +) -> ClassT | Callable[[ClassT], ClassT]: + """Mark a class as a Stream definition. Its :py:func:`topic` attributes are the Topics. + + A definition has no backend in it, so Workflows, Activities and clients share it. + + Args: + cls: The class. + name: The default name of a Stream of this definition. Defaults to the class name. + + .. warning:: + This API is experimental. + """ + + def decorate(cls: ClassT) -> ClassT: + # A subclass attribute replaces the base class attribute of the same name. + by_attribute: dict[str, TopicDefinition[Any]] = {} + for klass in reversed(cls.__mro__): + for attribute, value in vars(klass).items(): + if isinstance(value, TopicDefinition): + by_attribute[attribute] = value + topics: dict[str, TopicDefinition[Any]] = {} + for value in by_attribute.values(): + if value.name in topics: + raise ValueError( + f"The Stream definition {cls.__qualname__} has two Topics " + f"named {value.name}" + ) + topics[value.name] = value + setattr( + cls, + _DEFINITION_ATTR, + StreamDefinition(name=name or cls.__name__, topics=topics), + ) + return cls + + if cls is not None: + return decorate(cls) + return decorate + + +def definition_of(cls: type) -> StreamDefinition: + """The definition that :py:func:`defn` gave a class. + + Raises: + TypeError: The class has no ``@streams.defn``. + """ + definition = cls.__dict__.get(_DEFINITION_ATTR) if isinstance(cls, type) else None + if not isinstance(definition, StreamDefinition): + raise TypeError(f"{cls!r} is not a Stream definition. Add @streams.defn.") + return definition diff --git a/temporalio/contrib/streams/_handle.py b/temporalio/contrib/streams/_handle.py new file mode 100644 index 0000000000..2c8219b31b --- /dev/null +++ b/temporalio/contrib/streams/_handle.py @@ -0,0 +1,250 @@ +"""The shared part of the Stream handles of each context.""" + +from __future__ import annotations + +import typing +from typing import Any, Generic, TypeVar + +from typing_extensions import Self + +import temporalio.converter +from temporalio.contrib.streams import _errors +from temporalio.contrib.streams._defn import TopicDefinition, definition_of + +E = TypeVar("E") +H = TypeVar("H", bound="StreamHandle[Any]") + + +def _to_json(handle: StreamHandle[Any]) -> dict[str, str]: + # The proto JSON of temporal.api.streamservice.v1.StreamRef, with empty fields left out as + # MessageToJson does, so every SDK reads a handle that another wrote. + json = {"id": handle.id} + if handle.workflow_id: + json["workflowId"] = handle.workflow_id + if handle.first_execution_run_id: + json["firstExecutionRunId"] = handle.first_execution_run_id + if handle.backend: + json["backend"] = handle.backend + return json + + +def _from_json(value: Any, type_hint: Any) -> StreamHandle[Any]: + origin = typing.get_origin(type_hint) or type_hint + if not isinstance(origin, type) or not issubclass(origin, StreamHandle): + raise TypeError(f"{type_hint!r} is not a Stream handle type") + if isinstance(value, origin): + return value + if isinstance(value, StreamHandle): + value = _to_json(value) + if not isinstance(value, dict) or not isinstance(value.get("id"), str): + raise TypeError(f"Expected the JSON object of a Stream, got {type(value)}") + if origin is StreamHandle: + raise TypeError( + "Give a Stream parameter the handle type of its context, such as " + "ActivityStream[Events]" + ) + args = typing.get_args(type_hint) + # Any is a class from Python 3.11, and a hint such as ActivityStream[Any] has no definition. + definition = ( + args[0] if args and args[0] is not Any and isinstance(args[0], type) else None + ) + return origin( + id=value["id"], + workflow_id=value.get("workflowId") or None, + first_execution_run_id=value.get("firstExecutionRunId") or None, + backend=value.get("backend") or None, + definition=definition, + ) + + +class _StreamHandleJson( + temporalio.converter.TransferTypeConverter["StreamHandle[Any]", dict[str, str]] +): + # One form for every handle type, and the parameter's type hint picks the class, so a + # Workflow passes a Stream to an Activity as data. + transfer_type: type[dict[str, str]] | None = dict + + def to_transfer_type(self, value: StreamHandle[Any]) -> dict[str, str]: + return _to_json(value) + + def from_transfer_type( + self, value: dict[str, str], type_hint: type[StreamHandle[Any]] + ) -> StreamHandle[Any]: + return _from_json(value, type_hint) + + +class StreamHandle(Generic[E]): + """Base of the Stream handle of each context. + + A handle is data: it passes as a Workflow or Activity argument, as the proto JSON of the + ``StreamRef`` message. The parameter's type, such as ``ActivityStream[Events]``, gives the + handle class and the definition. + + .. warning:: + This API is experimental. + """ + + def __init__( + self, + *, + id: str, + workflow_id: str | None = None, + first_execution_run_id: str | None = None, + backend: str | None = None, + definition: type[E] | None = None, + ) -> None: + """Create a handle. Use the create and get calls of a context instead.""" + self._id = id + self._workflow_id = workflow_id + self._first_execution_run_id = first_execution_run_id + self._backend = backend + self._definition = definition + if definition is not None: + definition_of(definition) + + @classmethod + def from_handle(cls: type[H], handle: StreamHandle[Any]) -> H: + """The same Stream as a handle of this class, such as an ``ActivityStream`` from a + ``WorkflowStream``. + + A handle passes as data, so a Workflow uses it to give an Activity's input the type + that the Activity declares. + """ + return cls( + id=handle.id, + workflow_id=handle.workflow_id, + first_execution_run_id=handle.first_execution_run_id, + backend=handle.backend, + definition=handle.definition, + ) + + @property + def id(self) -> str: + """The Stream id. It is opaque to callers.""" + return self._id + + @property + def workflow_id(self) -> str | None: + """For a Workflow's Stream, the id of the Workflow that owns it.""" + return self._workflow_id + + @property + def first_execution_run_id(self) -> str | None: + """For a Workflow's Stream, the first run id of the owning run chain.""" + return self._first_execution_run_id + + @property + def backend(self) -> str | None: + """Name of the backend that keeps the Stream. None means the default one.""" + return self._backend + + @property + def definition(self) -> type[E] | None: + """The Stream definition class, or None for a Stream with no definition.""" + return self._definition + + def _topic_spec( + self, topic: str | TopicDefinition[Any], type: Any + ) -> tuple[str, Any]: + # The name and the decode type of a Topic. A Topic of a definition keeps its type, + # and a name outside the definition needs a type, so a typo does not pass. + definition = definition_of(self._definition) if self._definition else None + if isinstance(topic, TopicDefinition): + if type is not None: + raise TypeError("A Topic of a definition already has its type") + if ( + definition is not None + and definition.topics.get(topic.name) is not topic + ): + raise _errors.StreamError( + f"The Stream definition {definition.name} has no Topic {topic.name}", + reason="unknown_topic", + ) + return topic.name, topic.type + if definition is not None: + known = definition.topics.get(topic) + if known is not None: + if type is not None and type is not known.type: + raise TypeError( + f"The Topic {topic} of {definition.name} has the type " + f"{known.type!r}" + ) + return topic, known.type + if type is None: + raise _errors.StreamError( + f"The Stream definition {definition.name} has no Topic {topic}. " + "Pass type= for a Topic outside the definition.", + reason="unknown_topic", + ) + return topic, type + + def _known_types(self) -> dict[str, Any]: + # A read of every Topic still decodes each Topic of the definition with its type. + if self._definition is None: + return {} + topics = definition_of(self._definition).topics + return {name: topic.type for name, topic in topics.items()} + + def _fields(self) -> dict[str, str]: + # The fields that name the Stream. They need no call, so a handle whose id is not known + # yet still compares, hashes and prints. + fields = {"id": self._id} + if self._workflow_id: + fields["workflowId"] = self._workflow_id + if self._first_execution_run_id: + fields["firstExecutionRunId"] = self._first_execution_run_id + if self._backend: + fields["backend"] = self._backend + return fields + + def __eq__(self, other: object) -> bool: + """Handles of one class are equal when they name the same Stream.""" + if type(other) is not type(self): + return NotImplemented + assert isinstance(other, StreamHandle) + return self._fields() == other._fields() + + def __hash__(self) -> int: + """Hash of the fields that never change, so it stays when a client learns the id.""" + return hash((type(self), self._backend, self._workflow_id or self._id)) + + def __repr__(self) -> str: + """Show the Stream fields.""" + fields = ", ".join(f"{key}={value!r}" for key, value in self._fields().items()) + return f"{type(self).__name__}({fields})" + + def __copy__(self) -> Self: + """A copy is a handle of the same Stream in the same context.""" + return self + + def __deepcopy__(self, memo: dict[int, Any]) -> Self: + """A deep copy is the same handle, since dataclasses.asdict copies its fields.""" + return self + + @classmethod + def __get_pydantic_core_schema__(cls, source: Any, handler: Any) -> Any: + """Give pydantic models the proto JSON form, as the payload converter does. + + Pydantic is optional, so it is imported only when a model uses a handle. + """ + from pydantic_core import core_schema + + return core_schema.no_info_plain_validator_function( + lambda value: _from_json(value, source), + json_schema_input_schema=core_schema.typed_dict_schema( + { + "id": core_schema.typed_dict_field(core_schema.str_schema()), + **{ + key: core_schema.typed_dict_field( + core_schema.str_schema(), required=False + ) + for key in ("workflowId", "firstExecutionRunId", "backend") + }, + } + ), + serialization=core_schema.plain_serializer_function_ser_schema(_to_json), + ) + + +# Registered as a call, not a decorator, so type checkers keep the class generic. +temporalio.converter.transfer_type_convertible(_StreamHandleJson)(StreamHandle) diff --git a/temporalio/converter/_payload_converter.py b/temporalio/converter/_payload_converter.py index 8b428c2290..3540d743cb 100644 --- a/temporalio/converter/_payload_converter.py +++ b/temporalio/converter/_payload_converter.py @@ -110,6 +110,10 @@ def transfer_type_convertible( ) -> _TransferTypeConvertibleDecorator[ValueT, TransferTypeT]: """Decorate a class with a transfer type converter class. + A value of the class takes its transfer type at the top level of a payload, and inside a + JSON value such as a list or a dataclass field. A dataclass is the exception inside a + JSON value: ``dataclasses.asdict`` writes it as its fields, and it decodes from them. + .. warning:: This API is experimental and subject to change. """ @@ -125,6 +129,17 @@ def _get_transfer_type_converter( return None +def _nested_transfer_type_converter( + hint: object, +) -> TransferTypeConverter[Any, Any] | None: + # A dataclass inside another value is written as its fields, so only other classes take + # their transfer type there. + origin = typing.get_origin(hint) or hint + if not inspect.isclass(origin) or dataclasses.is_dataclass(origin): + return None + return _get_transfer_type_converter(origin) + + class PayloadConverter(ABC): """Base payload converter to/from multiple payloads/values.""" @@ -674,6 +689,14 @@ def default(self, o: Any) -> Any: # Datetime support if isinstance(o, datetime): return o.isoformat() + # A transfer type value inside another value. A dataclass never gets here, since + # dataclasses.asdict has already turned it into a dict. + converter = _get_transfer_type_converter(type(o)) + if converter is not None: + value = converter.to_transfer_type(o) + if isinstance(value, google.protobuf.message.Message): + return google.protobuf.json_format.MessageToDict(value) + return value # Dataclass support if dataclasses.is_dataclass(o) and not isinstance(o, type): return dataclasses.asdict(o) @@ -843,6 +866,20 @@ def value_to_type( if ret is not JSONTypeConverter.Unhandled: return ret + # A transfer type value inside another value, as AdvancedJSONEncoder writes it + converter = _nested_transfer_type_converter(hint) + if converter is not None: + transfer_type = converter.transfer_type + if inspect.isclass(transfer_type) and issubclass( + transfer_type, google.protobuf.message.Message + ): + value = google.protobuf.json_format.ParseDict( + value, transfer_type(), ignore_unknown_fields=True + ) + elif transfer_type is not None: + value = value_to_type(transfer_type, value, custom_converters) + return converter.from_transfer_type(value, hint) + # Any or primitives if hint is Any: return value diff --git a/tests/contrib/streams/test_handles.py b/tests/contrib/streams/test_handles.py new file mode 100644 index 0000000000..303526fe2f --- /dev/null +++ b/tests/contrib/streams/test_handles.py @@ -0,0 +1,180 @@ +import json +from dataclasses import dataclass +from typing import Any, TypeVar + +import pydantic +import pytest + +from temporalio.common import RawValue +from temporalio.contrib import streams +from temporalio.contrib.pydantic import pydantic_data_converter +from temporalio.contrib.streams._defn import definition_of +from temporalio.converter import DataConverter + +E = TypeVar("E") + + +@dataclass +class Turn: + kind: str + + +@streams.defn(name="events") +class Events: + turns = streams.topic(Turn) + tokens = streams.topic(str) + raw = streams.topic(RawValue, name="raw-values") + + +@streams.defn +class Index: + turns = streams.topic(int) + + +class PlainStream(streams.StreamHandle[E]): + pass + + +OWNED = PlainStream( + id="wf:run:events", + workflow_id="wf", + first_execution_run_id="run", + backend="redis", + definition=Events, +) + + +def test_definition_takes_the_name_and_the_topics(): + events = definition_of(Events) + assert events.name == "events" + assert list(events.topics) == ["turns", "tokens", "raw-values"] + assert Events.turns.name == "turns" and Events.turns.type is Turn + assert Events.raw.name == "raw-values" + assert definition_of(Index).name == "Index" + + +def test_definition_rejects_two_topics_of_one_name(): + with pytest.raises(ValueError, match="two Topics named t"): + + @streams.defn + class Twice: # pyright: ignore[reportUnusedClass] + a = streams.topic(int, name="t") + b = streams.topic(str, name="t") + + +def test_subclass_topic_replaces_the_base_topic(): + class Base: + turns = streams.topic(int) + + @streams.defn + class Sub(Base): + turns = streams.topic(str) + + assert definition_of(Sub).topics["turns"].type is str + + +def test_class_without_defn_is_not_a_definition(): + with pytest.raises(TypeError, match="not a Stream definition"): + PlainStream(id="s", definition=Turn) + + +def test_topic_of_the_definition_keeps_its_type(): + assert OWNED._topic_spec(Events.turns, None) == ("turns", Turn) + assert OWNED._topic_spec("tokens", None) == ("tokens", str) + assert OWNED._topic_spec("custom", Turn) == ("custom", Turn) + with pytest.raises(streams.StreamError, match="has no Topic turns") as err: + OWNED._topic_spec(Index.turns, None) + assert err.value.reason == "unknown_topic" + with pytest.raises(streams.StreamError, match="Pass type="): + OWNED._topic_spec("custom", None) + with pytest.raises(TypeError, match="has the type"): + OWNED._topic_spec("tokens", int) + with pytest.raises(TypeError, match="already has its type"): + OWNED._topic_spec(Events.turns, int) + + +def test_stream_without_definition_takes_any_topic(): + plain: PlainStream[Any] = PlainStream(id="s") + assert plain._topic_spec("t", None) == ("t", None) + assert plain._topic_spec(Events.turns, None) == ("turns", Turn) + + +def test_handle_is_the_proto_json_of_the_ref(): + converter = DataConverter.default.payload_converter + payload = converter.to_payload(OWNED) + assert json.loads(payload.data) == { + "id": "wf:run:events", + "workflowId": "wf", + "firstExecutionRunId": "run", + "backend": "redis", + } + assert json.loads(converter.to_payload(PlainStream(id="s")).data) == {"id": "s"} + + +def test_type_hint_picks_the_class_and_the_definition(): + converter = DataConverter.default.payload_converter + payload = converter.to_payload(OWNED) + decoded = converter.from_payload(payload, PlainStream[Events]) + assert type(decoded) is PlainStream + assert decoded == OWNED + assert decoded.definition is Events + bare = converter.from_payload(payload, PlainStream) + assert bare.definition is None and bare.backend == "redis" + with pytest.raises(Exception, match="handle type of its context"): + converter.from_payload(payload, streams.StreamHandle) + + +@dataclass +class CallModelInput: + prompt: str + stream: PlainStream[Events] + + +def test_handle_in_a_dataclass_field(): + converter = DataConverter.default.payload_converter + payload = converter.to_payload(CallModelInput("hi", OWNED)) + assert json.loads(payload.data)["stream"]["workflowId"] == "wf" + decoded = converter.from_payload(payload, CallModelInput) + assert decoded.stream == OWNED and decoded.stream.definition is Events + + +class CallModelModel(pydantic.BaseModel): + prompt: str + stream: PlainStream[Events] + + +def test_handle_in_a_pydantic_field(): + model = CallModelModel(prompt="hi", stream=OWNED) + dumped = json.loads(model.model_dump_json()) + assert dumped["stream"] == { + "id": "wf:run:events", + "workflowId": "wf", + "firstExecutionRunId": "run", + "backend": "redis", + } + loaded = CallModelModel.model_validate_json(model.model_dump_json()) + assert loaded.stream == OWNED and loaded.stream.definition is Events + converter = pydantic_data_converter.payload_converter + payload = converter.to_payload(model) + decoded = converter.from_payload(payload, CallModelModel) + assert decoded.stream == OWNED + + +class OtherStream(streams.StreamHandle[E]): + pass + + +def test_from_handle_keeps_the_stream(): + other: OtherStream[Events] = OtherStream.from_handle(OWNED) + assert type(other) is OtherStream + assert other.definition is Events + assert DataConverter.default.payload_converter.to_payload( + other + ) == DataConverter.default.payload_converter.to_payload(OWNED) + + +def test_any_type_argument_has_no_definition(): + converter = DataConverter.default.payload_converter + decoded = converter.from_payload(converter.to_payload(OWNED), PlainStream[Any]) + assert type(decoded) is PlainStream and decoded.definition is None + assert decoded.id == OWNED.id diff --git a/tests/test_converter.py b/tests/test_converter.py index 3237b71a75..d21d9193ab 100644 --- a/tests/test_converter.py +++ b/tests/test_converter.py @@ -3,6 +3,7 @@ import dataclasses import inspect import ipaddress +import json import logging import sys import traceback @@ -461,6 +462,56 @@ def test_temporal_transfer_type_payload_converter_with_generic_value( assert converter.from_payload(payload, type_hint) == value +class NestedTransferTypeValueConverter( + TransferTypeConverter["NestedTransferTypeValue[Any]", dict[str, str]] +): + transfer_type: type[dict[str, str]] | None = dict + + def to_transfer_type(self, value: NestedTransferTypeValue[Any]) -> dict[str, str]: + return {"wireName": value.name} + + def from_transfer_type( + self, value: dict[str, str], type_hint: type[NestedTransferTypeValue[Any]] + ) -> NestedTransferTypeValue[Any]: + args = typing.get_args(type_hint) + return NestedTransferTypeValue(value["wireName"], args[0] if args else None) + + +class NestedTransferTypeValue(Generic[T]): + def __init__(self, name: str, kind: Any = None) -> None: + self.name = name + self.kind = kind + + +transfer_type_convertible(NestedTransferTypeValueConverter)(NestedTransferTypeValue) + + +@dataclass +class HasNestedTransferTypeValues: + one: NestedTransferTypeValue[int] + many: list[NestedTransferTypeValue[str]] + proto: TemporalTransferTypeValueWithoutHint + + +def test_transfer_type_values_nested_in_json(): + converter = DataConverter.default.payload_converter + value = HasNestedTransferTypeValues( + one=NestedTransferTypeValue("a"), + many=[NestedTransferTypeValue("b")], + proto=TemporalTransferTypeValueWithoutHint("c"), + ) + payload = converter.to_payload(value) + assert json.loads(payload.data) == { + "one": {"wireName": "a"}, + "many": [{"wireName": "b"}], + "proto": {"value": "c"}, + } + decoded = converter.from_payload(payload, HasNestedTransferTypeValues) + assert decoded.one.name == "a" and decoded.one.kind is int + assert decoded.many[0].name == "b" and decoded.many[0].kind is str + assert decoded.proto == TemporalTransferTypeValueWithoutHint("c") + + def test_transfer_type_convertible_rejects_existing_converter(): with pytest.raises(TypeError, match="already has a transfer type converter"): transfer_type_convertible(TemporalTransferTypeValueConverter)(