diff --git a/GUIDE.md b/GUIDE.md index c12b922b7..0f5d61bb6 100644 --- a/GUIDE.md +++ b/GUIDE.md @@ -2303,7 +2303,7 @@ type or transform expression refers to an otherwise-unavailable module. ### @nexus.poll **Placement:** Operation (function) **Syntax:** -`@nexus.poll python="" [target="workflow|client"]` +`@nexus.poll python="" [target="workflow|client"] [items=""] [python-resume=""]` Generates an `iter_` async iterator for the operation, and for each resource method that binds it. The iterator sends the request, yields the @@ -2338,6 +2338,28 @@ def next_read_request(request, response): The referenced helper must be available from the package specified by `--support-package`. +`items` names a list field of the response. The iterator then yields the items, +not the pages, and it is a `PollIterator` of the package. Its `next_request` is +the request that resumes after the last item that it yielded, and it is `None` +after the last page. `resume` names a helper `(page request, item) -> request` +that resumes after one item. After a page's last item, or after a page with no +item, `next_request` is the page helper's request. A caller saves it and resumes +with it. + +```wit +/// @nexus.poll python="next_read_request" items="records" +/// python-resume="resume_read_request" target="client" +read-records: func(request: read-records-request) -> read-records-response; +``` + +```python +def resume_read_request(request, item): + return dataclasses.replace(request, from_offset=(item.offset or 0) + 1, last=None) +``` + +nexgen refuses `items` that names no list field of the response, and `resume` +with no `items`. + --- ### @nexus.flatten-in-api diff --git a/advanced/samples/README.md b/advanced/samples/README.md index 27b513418..380366a13 100644 --- a/advanced/samples/README.md +++ b/advanced/samples/README.md @@ -75,7 +75,7 @@ sample-project instructions for [Go](go/README.md), [Python](python/README.md), Operations with `@nexus.poll`, in a service with an endpoint and in a service with no endpoint, and a resource method that binds one. Only Python generates the iterators, for both targets. The service with no endpoint also takes a -parameter of an enum type. +parameter of an enum type, and its poll yields the executions of each page. - WIT: [`inputs/poll-service.wit`](inputs/poll-service.wit) - Python: [`python/wit/poll_service/`](python/wit/poll_service/), @@ -213,7 +213,8 @@ each target, which the sample support code gives in place of Brook's runtime. The create takes the backend of the Stream's ref, and the ref copies the id. In Python, the factory `stream` builds a handle from an id or a ref with no call, and a handle passes as its ref. `publish` takes the values, and returns the -offset of the last one. +offset of the last one. The read iterator yields each Record, and its +`next_request` resumes after the last one. - WIT: [`inputs/stream-service.wit`](inputs/stream-service.wit), [`inputs/stream-writer.wit`](inputs/stream-writer.wit) diff --git a/advanced/samples/inputs/poll-service.wit b/advanced/samples/inputs/poll-service.wit index 1231e8b91..d4dd352bb 100644 --- a/advanced/samples/inputs/poll-service.wit +++ b/advanced/samples/inputs/poll-service.wit @@ -17,10 +17,69 @@ interface page-service { query: option, } + /// @nexus.proto "temporal.api.common.v1.WorkflowExecution" + record workflow-execution { + workflow-id: string, + run-id: string, + } + + /// @nexus.proto "temporal.api.workflow.v1.WorkflowExecutionInfo" + record workflow-execution-info { + execution: workflow-execution, + parent-execution: option, + /// @nexus.omit + %type: placeholder, + /// @nexus.omit + start-time: placeholder, + /// @nexus.omit + close-time: placeholder, + /// @nexus.omit + status: placeholder, + /// @nexus.omit + history-length: placeholder, + /// @nexus.omit + parent-namespace-id: placeholder, + /// @nexus.omit + execution-time: placeholder, + /// @nexus.omit + memo: placeholder, + /// @nexus.omit + search-attributes: placeholder, + /// @nexus.omit + auto-reset-points: placeholder, + /// @nexus.omit + task-queue: placeholder, + /// @nexus.omit + state-transition-count: placeholder, + /// @nexus.omit + history-size-bytes: placeholder, + /// @nexus.omit + most-recent-worker-version-stamp: placeholder, + /// @nexus.omit + execution-duration: placeholder, + /// @nexus.omit + root-execution: placeholder, + /// @nexus.omit + assigned-build-id: placeholder, + /// @nexus.omit + inherited-build-id: placeholder, + /// @nexus.omit + first-run-id: placeholder, + /// @nexus.omit + versioning-info: placeholder, + /// @nexus.omit + worker-deployment-name: placeholder, + /// @nexus.omit + priority: placeholder, + /// @nexus.omit + external-payload-size-bytes: placeholder, + /// @nexus.omit + external-payload-count: placeholder, + } + /// @nexus.proto "temporal.api.workflowservice.v1.ListWorkflowExecutionsResponse" record list-workflows-response { - /// @nexus.omit - executions: placeholder, + executions: list, next-page-token: option>, } @@ -58,6 +117,6 @@ interface feed-service { open-feed: func(request: open-feed-request) -> feed-result; - /// @nexus.poll python="next_page_request" + /// @nexus.poll python="next_page_request" items="executions" list-feed: func(request: list-workflows-request) -> list-workflows-response; } diff --git a/advanced/samples/inputs/stream-service.wit b/advanced/samples/inputs/stream-service.wit index 9e10d5404..7713dc593 100644 --- a/advanced/samples/inputs/stream-service.wit +++ b/advanced/samples/inputs/stream-service.wit @@ -428,7 +428,8 @@ interface stream-service { /// @nexus.doc /// "A long poll. NOT_FOUND for an uncreated Stream. Workflow code reads from a Local Activity." /// @nexus.operation name="ReadRecords" - /// @nexus.poll python="next_read_request" target="client" + /// @nexus.poll python="next_read_request" items="records" python-resume="resume_read_request" + /// target="client" read-records: func(request: read-records-request) -> read-records-response; /// @nexus.operation name="DescribeStream" diff --git a/advanced/samples/invalid/poll/invalid-resume-helper/expected.toml b/advanced/samples/invalid/poll/invalid-resume-helper/expected.toml new file mode 100644 index 000000000..15f8f1088 --- /dev/null +++ b/advanced/samples/invalid/poll/invalid-resume-helper/expected.toml @@ -0,0 +1,3 @@ +[error] +kind = "invalid-wit-directive" +contains = ["@nexus.poll", "invalid `resume` support helper `resume_read(request, item)`"] diff --git a/advanced/samples/invalid/poll/invalid-resume-helper/input.wit b/advanced/samples/invalid/poll/invalid-resume-helper/input.wit new file mode 100644 index 000000000..e07aa58e4 --- /dev/null +++ b/advanced/samples/invalid/poll/invalid-resume-helper/input.wit @@ -0,0 +1,14 @@ +package temporal:poll-invalid-resume-helper@1.0.0; + +world service { export reader; } + +/// @nexus.endpoint "reader" +interface reader { + record read-request { offset: option } + record read-item { offset: option } + record read-response { items: list, done: option } + + // The resume helper names a support helper, not a call. + /// @nexus.poll python="next_read" items="items" python-resume="resume_read(request, item)" + read: func(request: read-request) -> read-response; +} diff --git a/advanced/samples/invalid/poll/items-not-a-list/expected.toml b/advanced/samples/invalid/poll/items-not-a-list/expected.toml new file mode 100644 index 000000000..821282ddc --- /dev/null +++ b/advanced/samples/invalid/poll/items-not-a-list/expected.toml @@ -0,0 +1,3 @@ +[error] +kind = "invalid-wit-directive" +contains = ["@nexus.poll", "`done` isn't a list"] diff --git a/advanced/samples/invalid/poll/items-not-a-list/input.wit b/advanced/samples/invalid/poll/items-not-a-list/input.wit new file mode 100644 index 000000000..deca22d8d --- /dev/null +++ b/advanced/samples/invalid/poll/items-not-a-list/input.wit @@ -0,0 +1,14 @@ +package temporal:poll-items-not-a-list@1.0.0; + +world service { export reader; } + +/// @nexus.endpoint "reader" +interface reader { + record read-request { offset: option } + record read-item { offset: option } + record read-response { items: list, done: option } + + // A poll yields the items of a list. + /// @nexus.poll python="next_read" items="done" + read: func(request: read-request) -> read-response; +} diff --git a/advanced/samples/invalid/poll/items-unknown-field/expected.toml b/advanced/samples/invalid/poll/items-unknown-field/expected.toml new file mode 100644 index 000000000..0be11d92e --- /dev/null +++ b/advanced/samples/invalid/poll/items-unknown-field/expected.toml @@ -0,0 +1,3 @@ +[error] +kind = "invalid-wit-directive" +contains = ["@nexus.poll", "`records` isn't a field of the response"] diff --git a/advanced/samples/invalid/poll/items-unknown-field/input.wit b/advanced/samples/invalid/poll/items-unknown-field/input.wit new file mode 100644 index 000000000..e6cf356e7 --- /dev/null +++ b/advanced/samples/invalid/poll/items-unknown-field/input.wit @@ -0,0 +1,14 @@ +package temporal:poll-items-unknown-field@1.0.0; + +world service { export reader; } + +/// @nexus.endpoint "reader" +interface reader { + record read-request { offset: option } + record read-item { offset: option } + record read-response { items: list, done: option } + + // `records` isn't a field of the response. + /// @nexus.poll python="next_read" items="records" + read: func(request: read-request) -> read-response; +} diff --git a/advanced/samples/invalid/poll/resume-without-items/expected.toml b/advanced/samples/invalid/poll/resume-without-items/expected.toml new file mode 100644 index 000000000..40694a165 --- /dev/null +++ b/advanced/samples/invalid/poll/resume-without-items/expected.toml @@ -0,0 +1,3 @@ +[error] +kind = "invalid-wit-directive" +contains = ["@nexus.poll", "`resume` needs `items`"] diff --git a/advanced/samples/invalid/poll/resume-without-items/input.wit b/advanced/samples/invalid/poll/resume-without-items/input.wit new file mode 100644 index 000000000..9b199aaaf --- /dev/null +++ b/advanced/samples/invalid/poll/resume-without-items/input.wit @@ -0,0 +1,14 @@ +package temporal:poll-resume-without-items@1.0.0; + +world service { export reader; } + +/// @nexus.endpoint "reader" +interface reader { + record read-request { offset: option } + record read-item { offset: option } + record read-response { items: list, done: option } + + // A resume point is after an item, so it needs the items. + /// @nexus.poll python="next_read" python-resume="resume_read" + read: func(request: read-request) -> read-response; +} diff --git a/advanced/samples/python/temporal_support/__init__.py b/advanced/samples/python/temporal_support/__init__.py index 181e22d26..60026e62b 100644 --- a/advanced/samples/python/temporal_support/__init__.py +++ b/advanced/samples/python/temporal_support/__init__.py @@ -5,5 +5,6 @@ from .operations import * # noqa: F401,F403 from .poll import next_page_request as next_page_request from .poll import next_read_request as next_read_request +from .poll import resume_read_request as resume_read_request from .temporal_model_converters import * # noqa: F401,F403 from .streams import * # noqa: F401,F403 diff --git a/advanced/samples/python/temporal_support/poll.py b/advanced/samples/python/temporal_support/poll.py index 35da5aec1..5769d6f8c 100644 --- a/advanced/samples/python/temporal_support/poll.py +++ b/advanced/samples/python/temporal_support/poll.py @@ -47,3 +47,13 @@ def next_page_request(request: _RequestT, response: _PageResponse) -> _RequestT if not response.next_page_token: return None return dataclasses.replace(request, next_page_token=response.next_page_token) + + +class _ReadItem(typing.Protocol): + @property + def offset(self) -> int | None: ... + + +def resume_read_request(request: _ReadRequestT, item: _ReadItem) -> _ReadRequestT: + # A read that resumes after an item starts past its offset, and reads no `last` again. + return dataclasses.replace(request, from_offset=(item.offset or 0) + 1, last=None) diff --git a/advanced/samples/python/tests/test_client_target.py b/advanced/samples/python/tests/test_client_target.py index 6ddebe6f7..318ea19fc 100644 --- a/advanced/samples/python/tests/test_client_target.py +++ b/advanced/samples/python/tests/test_client_target.py @@ -280,9 +280,9 @@ async def test_a_read_poll_keeps_its_offset_on_an_empty_page() -> None: ] client = FakeClient({"ReadRecords": list(pages)}) - read = [page async for page in client.iter_read_records(str, stream=ref, last=2)] + read = [item async for item in client.iter_read_records(str, stream=ref, last=2)] - assert read == pages + assert read == [] assert [call.input for call in client.calls] == [ ReadRecordsRequest(stream=ref, last=2), ReadRecordsRequest(stream=ref, from_offset=3), @@ -291,23 +291,40 @@ async def test_a_read_poll_keeps_its_offset_on_an_empty_page() -> None: ] -async def test_a_handle_reads_every_page() -> None: +def stream_record(body: str, offset: int) -> StreamRecord[str]: + return StreamRecord(item=StreamItemData(StreamData(body=body)), offset=offset) + + +async def test_a_handle_yields_items_and_resumes_after_each() -> None: ref = StreamRef(id="s1") pages = [ - ReadRecordsResponse[str](next_offset=2), - ReadRecordsResponse[str](done=True), + ReadRecordsResponse( + records=[stream_record("a", 10), stream_record("b", 12)], next_offset=15 + ), + ReadRecordsResponse[str](next_offset=15), + ReadRecordsResponse( + records=[stream_record("c", 20)], next_offset=21, done=True + ), ] client = FakeClient( {"CreateStream": [CreateStreamResult(stream=ref)], "ReadRecords": list(pages)} ) stream = await client.create_stream(Texts, id="s1") - read = [page async for page in stream.iter_read(from_offset=1)] + read = stream.iter_read(from_offset=10) + seen: list[tuple[int | None, int | None]] = [] + async for record in read: + seen.append( + (record.offset, read.next_request and read.next_request.from_offset) + ) - assert read == pages + # After an item, the read resumes past it, and after a page's last item, where the page ends. + assert seen == [(10, 11), (12, 15), (20, None)] + assert read.next_request is None assert [call.input for call in client.calls[1:]] == [ - ReadRecordsRequest(stream=ref, from_offset=1), - ReadRecordsRequest(stream=ref, from_offset=2), + ReadRecordsRequest(stream=ref, from_offset=10), + ReadRecordsRequest(stream=ref, from_offset=15), + ReadRecordsRequest(stream=ref, from_offset=15), ] diff --git a/advanced/samples/python/tests/test_poll_service.py b/advanced/samples/python/tests/test_poll_service.py index b047d6acf..0c2add0d0 100644 --- a/advanced/samples/python/tests/test_poll_service.py +++ b/advanced/samples/python/tests/test_poll_service.py @@ -23,14 +23,26 @@ ListWorkflowsRequest, ListWorkflowsResponse, OpenFeedRequest, + WorkflowExecution, + WorkflowExecutionInfo, ) -# Each page names the next one, and the last page names none. + +def info(workflow_id: str) -> WorkflowExecutionInfo: + return WorkflowExecutionInfo( + execution=WorkflowExecution(workflow_id=workflow_id, run_id="r1") + ) + + +# Each page names the next one, and the last page names none. The middle page is empty. PAGES = { - None: ListWorkflowsResponse(next_page_token=b"2"), + None: ListWorkflowsResponse( + executions=[info("w1"), info("w2")], next_page_token=b"2" + ), b"2": ListWorkflowsResponse(next_page_token=b"3"), - b"3": ListWorkflowsResponse(), + b"3": ListWorkflowsResponse(executions=[info("w3")]), } +ITEMS = [info("w1"), info("w2"), info("w3")] def page(request: ListWorkflowsRequest) -> ListWorkflowsResponse: @@ -102,9 +114,12 @@ async def test_a_client_polls_an_operation_until_the_helper_stops() -> None: async def test_a_client_polls_a_service_with_no_endpoint() -> None: client = FakeClient() - pages = [response async for response in client.iter_list_feed(namespace="default")] + read = client.iter_list_feed(namespace="default") + items = [item async for item in read] - assert pages == list(PAGES.values()) + # The feed's poll yields the executions of each page. + assert items == ITEMS + assert read.next_request is None assert client.tokens("ListFeed") == [None, b"2", b"3"] assert {endpoint for endpoint, _, _ in client.calls} == {None} @@ -115,9 +130,9 @@ async def test_a_client_resource_polls_its_method() -> None: namespace="default", query="q", order=FeedOrder.Oldest ) - pages = [response async for response in feed.iter_read(page_size=10)] + items = [item async for item in feed.iter_read(page_size=10)] - assert pages == list(PAGES.values()) + assert items == ITEMS assert client.calls[0][2] == OpenFeedRequest( namespace="default", query="q", order=FeedOrder.Oldest ) @@ -167,12 +182,12 @@ class PollWorkflow: async def run(self) -> list[int]: pages = [p async for p in poll_service.iter_list_workflows(namespace="default")] service = poll_service.FeedServiceClient("feeds") - feed_pages = [p async for p in service.iter_list_feed(namespace="default")] + feed_items = [i async for i in service.iter_list_feed(namespace="default")] feed = await service.open_feed( namespace="default", query="q", order=FeedOrder.Oldest ) - read_pages = [p async for p in feed.iter_read()] - return [len(pages), len(feed_pages), len(read_pages)] + read_items = [i async for i in feed.iter_read()] + return [len(pages), len(feed_items), len(read_items)] async def test_workflow_code_polls_each_kind_of_operation( diff --git a/advanced/samples/python/wit/poll_service/_poll.py b/advanced/samples/python/wit/poll_service/_poll.py new file mode 100644 index 000000000..379105e39 --- /dev/null +++ b/advanced/samples/python/wit/poll_service/_poll.py @@ -0,0 +1,65 @@ +# Generated by nexgen latest. DO NOT EDIT! + +from __future__ import annotations + +import collections.abc +import typing + + +_PollRequestT = typing.TypeVar("_PollRequestT") +_PollItemT = typing.TypeVar("_PollItemT") + + +@typing.final +class PollIterator(typing.Generic[_PollRequestT, _PollItemT]): + """Yields each item that a poll reads, page by page. + + ``next_request`` resumes after the last item that the iterator yielded. It is ``None`` + after the last page. + """ + + def __init__( + self, + request: _PollRequestT, + read: collections.abc.Callable[ + [_PollRequestT], + collections.abc.Awaitable[ + tuple[_PollRequestT | None, collections.abc.Sequence[_PollItemT]] + ], + ], + resume: collections.abc.Callable[[_PollRequestT, _PollItemT], _PollRequestT] + | None, + ) -> None: + self.next_request: _PollRequestT | None = request + self._read: collections.abc.Callable[ + [_PollRequestT], + collections.abc.Awaitable[ + tuple[_PollRequestT | None, collections.abc.Sequence[_PollItemT]] + ], + ] = read + self._resume: ( + collections.abc.Callable[[_PollRequestT, _PollItemT], _PollRequestT] | None + ) = resume + self._page: list[_PollItemT] = [] + self._page_request: _PollRequestT = request + self._page_next: _PollRequestT | None = None + + def __aiter__(self) -> PollIterator[_PollRequestT, _PollItemT]: + return self + + async def __anext__(self) -> _PollItemT: + while not self._page: + if self.next_request is None: + raise StopAsyncIteration + self._page_request = self.next_request + self._page_next, items = await self._read(self._page_request) + self._page = list(items) + if not self._page: + self.next_request = self._page_next + item = self._page.pop(0) + # After a page's last item, the next read starts where the page helper says. + if not self._page: + self.next_request = self._page_next + elif self._resume is not None: + self.next_request = self._resume(self._page_request, item) + return item diff --git a/advanced/samples/python/wit/poll_service/_resources/feed.py b/advanced/samples/python/wit/poll_service/_resources/feed.py index 2d3d8449f..895e87ec2 100644 --- a/advanced/samples/python/wit/poll_service/_resources/feed.py +++ b/advanced/samples/python/wit/poll_service/_resources/feed.py @@ -10,8 +10,10 @@ from ..models import ( ListWorkflowsRequest, ListWorkflowsResponse, + WorkflowExecutionInfo, ) import temporal_support as _support +from .._poll import PollIterator @dataclasses.dataclass @@ -48,30 +50,35 @@ async def read( request, rpc_metadata=rpc_metadata, rpc_timeout=rpc_timeout ) - async def iter_read( + def iter_read( self, page_size: int | None = None, *, rpc_metadata: collections.abc.Mapping[str, str | bytes] | None = None, rpc_timeout: datetime.timedelta | None = None, - ) -> collections.abc.AsyncIterator[ListWorkflowsResponse]: - """Yield each response of ``read``. The poll helper gives the next request, and the - iteration stops when it gives ``None``. + ) -> PollIterator[ListWorkflowsRequest, WorkflowExecutionInfo]: + """Yield each item that ``read`` reads, page by page. The iterator's + ``next_request`` resumes after the last item that it yielded. """ request = ListWorkflowsRequest( namespace=self.namespace, page_size=page_size, query=self.query, ) - while True: + + async def read( + request: ListWorkflowsRequest, + ) -> tuple[ + ListWorkflowsRequest | None, collections.abc.Sequence[WorkflowExecutionInfo] + ]: response = await self._ops()._feed_service_list_feed( request, rpc_metadata=rpc_metadata, rpc_timeout=rpc_timeout ) - yield response - next_request = _support.next_page_request(request, response) - if next_request is None: - return - request = next_request + return _support.next_page_request( + request, response + ), response.executions or [] + + return PollIterator(request, read, None) class _FeedServiceOperations(typing.Protocol): diff --git a/advanced/samples/python/wit/poll_service/client.py b/advanced/samples/python/wit/poll_service/client.py index e2495cdb5..7fe78bc6d 100644 --- a/advanced/samples/python/wit/poll_service/client.py +++ b/advanced/samples/python/wit/poll_service/client.py @@ -13,9 +13,11 @@ ListWorkflowsRequest, ListWorkflowsResponse, OpenFeedRequest, + WorkflowExecutionInfo, ) from ._resources import Feed import temporal_support as _support +from ._poll import PollIterator _OutputT = typing.TypeVar("_OutputT") @@ -263,7 +265,7 @@ async def list_feed( request, rpc_metadata=rpc_metadata, rpc_timeout=rpc_timeout ) - async def iter_list_feed( + def iter_list_feed( self, *, namespace: str, @@ -272,9 +274,9 @@ async def iter_list_feed( query: str | None = None, rpc_metadata: collections.abc.Mapping[str, str | bytes] | None = None, rpc_timeout: datetime.timedelta | None = None, - ) -> collections.abc.AsyncIterator[ListWorkflowsResponse]: - """Yield each response of ``list_feed``. The poll helper gives the next request, - and the iteration stops when it gives ``None``. + ) -> PollIterator[ListWorkflowsRequest, WorkflowExecutionInfo]: + """Yield each item that ``list_feed`` reads, page by page. The iterator's + ``next_request`` resumes after the last item that it yielded. """ request = ListWorkflowsRequest( namespace=namespace, @@ -282,15 +284,20 @@ async def iter_list_feed( next_page_token=next_page_token, query=query, ) - while True: + + async def read( + request: ListWorkflowsRequest, + ) -> tuple[ + ListWorkflowsRequest | None, collections.abc.Sequence[WorkflowExecutionInfo] + ]: response = await self._feed_service_list_feed( request, rpc_metadata=rpc_metadata, rpc_timeout=rpc_timeout ) - yield response - next_request = _support.next_page_request(request, response) - if next_request is None: - return - request = next_request + return _support.next_page_request( + request, response + ), response.executions or [] + + return PollIterator(request, read, None) def feed( self, diff --git a/advanced/samples/python/wit/poll_service/models.py b/advanced/samples/python/wit/poll_service/models.py index bfef018a0..31c0cca5c 100644 --- a/advanced/samples/python/wit/poll_service/models.py +++ b/advanced/samples/python/wit/poll_service/models.py @@ -5,6 +5,8 @@ import dataclasses import enum import typing_extensions +import temporalio.api.common.v1.message_pb2 +import temporalio.api.workflow.v1.message_pb2 import temporalio.api.workflowservice.v1.request_response_pb2 import temporalio.converter @@ -86,6 +88,9 @@ def to_transfer_type( @dataclasses.dataclass(slots=True) class ListWorkflowsResponse: + executions: list[WorkflowExecutionInfo] | None = dataclasses.field( + default_factory=list + ) next_page_token: bytes | None = None @@ -109,6 +114,12 @@ def from_transfer_type( type_hint: type[ListWorkflowsResponse], ) -> ListWorkflowsResponse: return ListWorkflowsResponse( + executions=[ + _WorkflowExecutionInfoTransferTypeConverter().from_transfer_type( + item, WorkflowExecutionInfo + ) + for item in value.executions + ], next_page_token=value.next_page_token if bool(value.next_page_token) else None, @@ -120,6 +131,14 @@ def to_transfer_type( value: ListWorkflowsResponse, ) -> temporalio.api.workflowservice.v1.request_response_pb2.ListWorkflowExecutionsResponse: message = temporalio.api.workflowservice.v1.request_response_pb2.ListWorkflowExecutionsResponse() + if value.executions: + for element in value.executions: + item = message.executions.add() + item.CopyFrom( + _WorkflowExecutionInfoTransferTypeConverter().to_transfer_type( + element + ) + ) if value.next_page_token is not None: message.next_page_token = value.next_page_token return message @@ -130,6 +149,113 @@ def to_transfer_type( )(ListWorkflowsResponse) # pyright: ignore[reportUnusedCallResult] +@dataclasses.dataclass(slots=True) +class WorkflowExecutionInfo: + execution: WorkflowExecution + parent_execution: WorkflowExecution | None = None + + +class _WorkflowExecutionInfoTransferTypeConverter( + temporalio.converter.TransferTypeConverter[ + WorkflowExecutionInfo, + temporalio.api.workflow.v1.message_pb2.WorkflowExecutionInfo, + ] +): + transfer_type: ( + type[temporalio.api.workflow.v1.message_pb2.WorkflowExecutionInfo] | None + ) = temporalio.api.workflow.v1.message_pb2.WorkflowExecutionInfo + + @typing_extensions.override + def from_transfer_type( + self, + value: temporalio.api.workflow.v1.message_pb2.WorkflowExecutionInfo, + type_hint: type[WorkflowExecutionInfo], + ) -> WorkflowExecutionInfo: + if not value.HasField("execution"): + raise ValueError("missing required field WorkflowExecutionInfo.execution") + execution = _WorkflowExecutionTransferTypeConverter().from_transfer_type( + value.execution, WorkflowExecution + ) + return WorkflowExecutionInfo( + execution=execution, + parent_execution=_WorkflowExecutionTransferTypeConverter().from_transfer_type( + value.parent_execution, WorkflowExecution + ) + if value.HasField("parent_execution") + else None, + ) + + @typing_extensions.override + def to_transfer_type( + self, + value: WorkflowExecutionInfo, + ) -> temporalio.api.workflow.v1.message_pb2.WorkflowExecutionInfo: + message = temporalio.api.workflow.v1.message_pb2.WorkflowExecutionInfo() + message.execution.CopyFrom( + _WorkflowExecutionTransferTypeConverter().to_transfer_type(value.execution) + ) + if value.parent_execution is not None: + message.parent_execution.CopyFrom( + _WorkflowExecutionTransferTypeConverter().to_transfer_type( + value.parent_execution + ) + ) + return message + + +temporalio.converter.transfer_type_convertible( + _WorkflowExecutionInfoTransferTypeConverter +)(WorkflowExecutionInfo) # pyright: ignore[reportUnusedCallResult] + + +@dataclasses.dataclass(slots=True) +class WorkflowExecution: + workflow_id: str + run_id: str + + +class _WorkflowExecutionTransferTypeConverter( + temporalio.converter.TransferTypeConverter[ + WorkflowExecution, temporalio.api.common.v1.message_pb2.WorkflowExecution + ] +): + transfer_type: ( + type[temporalio.api.common.v1.message_pb2.WorkflowExecution] | None + ) = temporalio.api.common.v1.message_pb2.WorkflowExecution + + @typing_extensions.override + def from_transfer_type( + self, + value: temporalio.api.common.v1.message_pb2.WorkflowExecution, + type_hint: type[WorkflowExecution], + ) -> WorkflowExecution: + if not value.workflow_id: + raise ValueError("missing required field WorkflowExecution.workflow_id") + workflow_id = value.workflow_id + if not value.run_id: + raise ValueError("missing required field WorkflowExecution.run_id") + run_id = value.run_id + return WorkflowExecution( + workflow_id=workflow_id, + run_id=run_id, + ) + + @typing_extensions.override + def to_transfer_type( + self, + value: WorkflowExecution, + ) -> temporalio.api.common.v1.message_pb2.WorkflowExecution: + message = temporalio.api.common.v1.message_pb2.WorkflowExecution() + message.workflow_id = value.workflow_id + message.run_id = value.run_id + return message + + +temporalio.converter.transfer_type_convertible(_WorkflowExecutionTransferTypeConverter)( + WorkflowExecution +) # pyright: ignore[reportUnusedCallResult] + + @dataclasses.dataclass(slots=True) class OpenFeedRequest: namespace: str diff --git a/advanced/samples/python/wit/poll_service/services.py b/advanced/samples/python/wit/poll_service/services.py index e167e4a16..caee80998 100644 --- a/advanced/samples/python/wit/poll_service/services.py +++ b/advanced/samples/python/wit/poll_service/services.py @@ -12,12 +12,14 @@ ListWorkflowsRequest, ListWorkflowsResponse, OpenFeedRequest, + WorkflowExecutionInfo, ) from ._resources.feed import ( Feed, bind_feed_client, ) import temporal_support as _support +from ._poll import PollIterator if typing.TYPE_CHECKING: from ._resources import Feed @@ -91,16 +93,16 @@ async def list_feed( output_type=ListWorkflowsResponse, ) - async def iter_list_feed( + def iter_list_feed( self, *, namespace: str, page_size: int | None = None, next_page_token: bytes | None = None, query: str | None = None, - ) -> collections.abc.AsyncIterator[ListWorkflowsResponse]: - """Yield each response of ``list_feed``. The poll helper gives the next request, - and the iteration stops when it gives ``None``. + ) -> PollIterator[ListWorkflowsRequest, WorkflowExecutionInfo]: + """Yield each item that ``list_feed`` reads, page by page. The iterator's + ``next_request`` resumes after the last item that it yielded. """ request = ListWorkflowsRequest( namespace=namespace, @@ -108,13 +110,18 @@ async def iter_list_feed( next_page_token=next_page_token, query=query, ) - while True: + + async def read( + request: ListWorkflowsRequest, + ) -> tuple[ + ListWorkflowsRequest | None, collections.abc.Sequence[WorkflowExecutionInfo] + ]: handle = await self._nexus_client.start_operation( operation="ListFeed", input=request, output_type=ListWorkflowsResponse ) response = await handle - yield response - next_request = _support.next_page_request(request, response) - if next_request is None: - return - request = next_request + return _support.next_page_request( + request, response + ), response.executions or [] + + return PollIterator(request, read, None) diff --git a/advanced/samples/python/wit/stream_service/_poll.py b/advanced/samples/python/wit/stream_service/_poll.py new file mode 100644 index 000000000..379105e39 --- /dev/null +++ b/advanced/samples/python/wit/stream_service/_poll.py @@ -0,0 +1,65 @@ +# Generated by nexgen latest. DO NOT EDIT! + +from __future__ import annotations + +import collections.abc +import typing + + +_PollRequestT = typing.TypeVar("_PollRequestT") +_PollItemT = typing.TypeVar("_PollItemT") + + +@typing.final +class PollIterator(typing.Generic[_PollRequestT, _PollItemT]): + """Yields each item that a poll reads, page by page. + + ``next_request`` resumes after the last item that the iterator yielded. It is ``None`` + after the last page. + """ + + def __init__( + self, + request: _PollRequestT, + read: collections.abc.Callable[ + [_PollRequestT], + collections.abc.Awaitable[ + tuple[_PollRequestT | None, collections.abc.Sequence[_PollItemT]] + ], + ], + resume: collections.abc.Callable[[_PollRequestT, _PollItemT], _PollRequestT] + | None, + ) -> None: + self.next_request: _PollRequestT | None = request + self._read: collections.abc.Callable[ + [_PollRequestT], + collections.abc.Awaitable[ + tuple[_PollRequestT | None, collections.abc.Sequence[_PollItemT]] + ], + ] = read + self._resume: ( + collections.abc.Callable[[_PollRequestT, _PollItemT], _PollRequestT] | None + ) = resume + self._page: list[_PollItemT] = [] + self._page_request: _PollRequestT = request + self._page_next: _PollRequestT | None = None + + def __aiter__(self) -> PollIterator[_PollRequestT, _PollItemT]: + return self + + async def __anext__(self) -> _PollItemT: + while not self._page: + if self.next_request is None: + raise StopAsyncIteration + self._page_request = self.next_request + self._page_next, items = await self._read(self._page_request) + self._page = list(items) + if not self._page: + self.next_request = self._page_next + item = self._page.pop(0) + # After a page's last item, the next read starts where the page helper says. + if not self._page: + self.next_request = self._page_next + elif self._resume is not None: + self.next_request = self._resume(self._page_request, item) + return item diff --git a/advanced/samples/python/wit/stream_service/_resources/stream.py b/advanced/samples/python/wit/stream_service/_resources/stream.py index 93f0e641b..d9a9ebd7b 100644 --- a/advanced/samples/python/wit/stream_service/_resources/stream.py +++ b/advanced/samples/python/wit/stream_service/_resources/stream.py @@ -22,11 +22,13 @@ ReadRecordsRequest, ReadRecordsResponse, StreamReadStart, + StreamRecord, StreamRef, StreamStatus, _StreamRefTransferTypeConverter, ) import temporal_support as _support +from .._poll import PollIterator if typing.TYPE_CHECKING: import temporal_support @@ -112,7 +114,7 @@ async def read( request, self.body_type, rpc_metadata=rpc_metadata, rpc_timeout=rpc_timeout ) - async def iter_read( + def iter_read( self, from_offset: int | None = None, start: StreamReadStart | None = None, @@ -122,9 +124,9 @@ async def iter_read( *, rpc_metadata: collections.abc.Mapping[str, str | bytes] | None = None, rpc_timeout: datetime.timedelta | None = None, - ) -> collections.abc.AsyncIterator[ReadRecordsResponse[BodyT]]: - """Yield each response of ``read``. The poll helper gives the next request, and the - iteration stops when it gives ``None``. + ) -> PollIterator[ReadRecordsRequest, StreamRecord[BodyT]]: + """Yield each item that ``read`` reads, page by page. The iterator's + ``next_request`` resumes after the last item that it yielded. """ request = ReadRecordsRequest( stream=self.stream, @@ -134,18 +136,21 @@ async def iter_read( wait=wait, create_wait=create_wait, ) - while True: + + async def read( + request: ReadRecordsRequest, + ) -> tuple[ + ReadRecordsRequest | None, collections.abc.Sequence[StreamRecord[BodyT]] + ]: response = await self._ops()._stream_service_read_records( request, self.body_type, rpc_metadata=rpc_metadata, rpc_timeout=rpc_timeout, ) - yield response - next_request = _support.next_read_request(request, response) - if next_request is None: - return - request = next_request + return _support.next_read_request(request, response), response.records or [] + + return PollIterator(request, read, _support.resume_read_request) async def describe( self, diff --git a/advanced/samples/python/wit/stream_service/client.py b/advanced/samples/python/wit/stream_service/client.py index 24361b33a..a9f27eab7 100644 --- a/advanced/samples/python/wit/stream_service/client.py +++ b/advanced/samples/python/wit/stream_service/client.py @@ -21,11 +21,13 @@ ReadRecordsRequest, ReadRecordsResponse, StreamReadStart, + StreamRecord, StreamRef, StreamStatus, ) from ._resources import Stream import temporal_support as _support +from ._poll import PollIterator if typing.TYPE_CHECKING: import temporal_support @@ -323,7 +325,7 @@ async def read_records( request, body_type, rpc_metadata=rpc_metadata, rpc_timeout=rpc_timeout ) - async def iter_read_records( + def iter_read_records( self, body_type: type[BodyT], *, @@ -337,9 +339,9 @@ async def iter_read_records( max_bytes: int | None = None, rpc_metadata: collections.abc.Mapping[str, str | bytes] | None = None, rpc_timeout: datetime.timedelta | None = None, - ) -> collections.abc.AsyncIterator[ReadRecordsResponse[BodyT]]: - """Yield each response of ``read_records``. The poll helper gives the next request, - and the iteration stops when it gives ``None``. + ) -> PollIterator[ReadRecordsRequest, StreamRecord[BodyT]]: + """Yield each item that ``read_records`` reads, page by page. The iterator's + ``next_request`` resumes after the last item that it yielded. """ request = ReadRecordsRequest( stream=stream, @@ -351,15 +353,18 @@ async def iter_read_records( last=last, max_bytes=max_bytes, ) - while True: + + async def read( + request: ReadRecordsRequest, + ) -> tuple[ + ReadRecordsRequest | None, collections.abc.Sequence[StreamRecord[BodyT]] + ]: response = await self._stream_service_read_records( request, body_type, rpc_metadata=rpc_metadata, rpc_timeout=rpc_timeout ) - yield response - next_request = _support.next_read_request(request, response) - if next_request is None: - return - request = next_request + return _support.next_read_request(request, response), response.records or [] + + return PollIterator(request, read, _support.resume_read_request) async def _stream_service_describe_stream( self, diff --git a/changelog/added/yielding-yellowhammer.md b/changelog/added/yielding-yellowhammer.md new file mode 100644 index 000000000..4f74f9b48 --- /dev/null +++ b/changelog/added/yielding-yellowhammer.md @@ -0,0 +1 @@ +`@nexus.poll items="" python-resume=""` makes a Python poll iterator yield the items of each page, with a `next_request` that resumes after the last item that it yielded. diff --git a/src/generator/python.rs b/src/generator/python.rs index 402782c19..4743842ba 100644 --- a/src/generator/python.rs +++ b/src/generator/python.rs @@ -711,6 +711,21 @@ impl<'a> ApiPlanner<'a> { GeneratedFileOrigin::fixed("generated Python resources package initializer"), )?; } + if services + .iter() + .flat_map(|service| &service.operations) + .any(|operation| operation.poll_items.is_some()) + { + let mut source = String::new(); + render_generated_file_header(&mut source); + source.push_str("\nimport collections.abc\nimport typing\n\n\n"); + source.push_str(PYTHON_POLL_ITERATOR); + files.insert( + "_poll.py", + source, + GeneratedFileOrigin::fixed("generated Python poll iterator"), + )?; + } if !services.is_empty() { files.insert( "services.py", @@ -899,6 +914,7 @@ impl<'a> ApiPlanner<'a> { render_named_python_import(&mut output, "._resources", &used_resource_names); } render_support_import(&mut output, &body, self.support_package, false); + render_poll_iterator_import(&mut output, &body, "._poll"); render_definition_type_checking_imports( &mut output, &body, @@ -1095,6 +1111,7 @@ impl<'a> ApiPlanner<'a> { render_named_python_import(&mut output, ".", &used_resource_names); } render_support_import(&mut output, &body, support_package, false); + render_poll_iterator_import(&mut output, &body, ".._poll"); render_definition_type_checking_imports(&mut output, &body, resource.definition.value()); if !body.is_empty() { output.push('\n'); @@ -1769,24 +1786,24 @@ impl<'a> ApiPlanner<'a> { return; }; let method_name = python_field_name(&method.name); - output.push_str("\n async def "); + let page_annotation = if targets.workflow { + temporalio_workflow_type_annotation(service, &operation.output_annotation) + } else { + operation.output_annotation.clone() + }; + let (keyword, annotation) = operation.poll_iterator_signature(&page_annotation); + output.push_str("\n "); + output.push_str(keyword); output.push_str(&poll_function_name(&method_name)); output.push_str("(\n"); self.render_resource_method_params(output, method, targets.client); - output.push_str(" ) -> collections.abc.AsyncIterator["); - if targets.workflow { - output.push_str(&temporalio_workflow_type_annotation( - service, - &operation.output_annotation, - )); - } else { - output.push_str(&operation.output_annotation); - } - output.push_str("]:\n"); + output.push_str(" ) -> "); + output.push_str(&annotation); + output.push_str(":\n"); render_python_docstring( output, " ", - Some(&poll_function_summary(&method_name)), + Some(&poll_function_summary(&method_name, operation)), &[], None, false, @@ -1814,7 +1831,7 @@ impl<'a> ApiPlanner<'a> { "response = await handle".to_string(), ] }; - render_poll_loop(output, " ", &response, next_expr); + render_poll_body(output, " ", &response, next_expr, operation); } fn python_resource_method_result_annotation(&self, method: &PlannedResourceMethod) -> String { @@ -2157,6 +2174,27 @@ impl<'a> ApiPlanner<'a> { }); } } + let poll_items = match (&operation.poll_items, operation.output_type()) { + (Some(key), Some(PlannedType::Record(record))) => { + let item = self + .api_plan + .record(&record.full_name) + .and_then(|record| record.fields.get(key)) + .and_then(|field| match &field.field_type { + PlannedType::List(item) => Some(item.as_ref()), + _ => None, + }); + match (item, self.model_field_attr(&record.full_name, key)) { + (Some(item), Some(attr)) => Some(( + attr, + resolve_python_value_type(self.api_plan, &self.external_models, item)? + .annotation, + )), + _ => None, + } + } + _ => None, + }; let descriptor_output_ref = match operation.output_type() { Some(PlannedType::Record(record)) => python_erased_generic_record_annotation( &output_ref, @@ -2201,6 +2239,11 @@ impl<'a> ApiPlanner<'a> { .map(|helper| format!("_support.{helper}")), poll_next_expr, poll_target: operation.poll_target, + poll_items, + poll_resume_expr: operation + .poll_resume + .value() + .map(|helper| format!("_support.{helper}")), output_resource_return, output_direct_result, output_none: operation.output_type().is_none(), @@ -3594,6 +3637,10 @@ struct RenderedOperation<'a> { /// The support helper that gives the next request of a poll. poll_next_expr: Option, poll_target: Option, + /// The response attribute whose items the poll yields, with the item's annotation. + poll_items: Option<(String, String)>, + /// The support helper that resumes a read after one item. + poll_resume_expr: Option, output_resource_return: Option, output_direct_result: bool, output_none: bool, @@ -3708,6 +3755,28 @@ impl RenderedOperation<'_> { } } + /// The keyword and the return annotation of a poll iterator. An item poll returns a + /// `PollIterator` at once, and another poll is an async generator of pages. + fn poll_iterator_signature(&self, page_annotation: &str) -> (&'static str, String) { + match &self.poll_items { + Some((_, item)) => ( + "def ", + format!("PollIterator[{}, {item}]", self.poll_request_annotation()), + ), + None => ( + "async def ", + format!("collections.abc.AsyncIterator[{page_annotation}]"), + ), + } + } + + fn poll_request_annotation(&self) -> &str { + self.unpacked_input + .as_ref() + .map(|input| input.model_name.as_str()) + .unwrap_or_else(|| operation_input_annotation(self)) + } + /// The poll helper, when the target gets the poll. fn poll_next(&self, target: GenerationTarget) -> Option<&str> { self.poll_next_expr @@ -5567,6 +5636,7 @@ fn render_service_module( } } render_support_import(&mut output, &endpoint_client_body, support_package, false); + render_poll_iterator_import(&mut output, &endpoint_client_body, "._poll"); render_definition_type_checking_imports( &mut output, &endpoint_client_body, @@ -5719,6 +5789,7 @@ fn render_operation_module( render_named_python_import(&mut output, ".._resources", &used_resource_names); } render_support_import(&mut output, &body, support_package, false); + render_poll_iterator_import(&mut output, &body, ".._poll"); render_definition_type_checking_imports( &mut output, &body, @@ -6332,7 +6403,11 @@ fn render_endpoint_service_object_poll_method( operation: &RenderedOperation<'_>, next_expr: &str, ) { - output.push_str(" async def "); + let (keyword, annotation) = operation.poll_iterator_signature( + &temporalio_workflow_type_annotation(service, &operation.output_annotation), + ); + output.push_str(" "); + output.push_str(keyword); output.push_str(&poll_function_name(&operation.attr_name)); output.push_str("(\n"); output.push_str(" self,\n"); @@ -6359,16 +6434,13 @@ fn render_endpoint_service_object_poll_method( output.push_str(operation_input_annotation(operation)); output.push_str(",\n"); } - output.push_str(" ) -> collections.abc.AsyncIterator["); - output.push_str(&temporalio_workflow_type_annotation( - service, - &operation.output_annotation, - )); - output.push_str("]:\n"); + output.push_str(" ) -> "); + output.push_str(&annotation); + output.push_str(":\n"); render_python_docstring( output, " ", - Some(&poll_function_summary(&operation.attr_name)), + Some(&poll_function_summary(&operation.attr_name, operation)), &[], None, false, @@ -6391,7 +6463,7 @@ fn render_endpoint_service_object_poll_method( nexus_client_start_line("self._nexus_client", operation, ""), "response = await handle".to_string(), ]; - render_poll_loop(output, " ", &response, next_expr); + render_poll_body(output, " ", &response, next_expr, operation); } fn render_endpoint_service_object_operation_method( @@ -6654,6 +6726,67 @@ fn validate_client_names(services: &[RenderedService<'_>]) -> Result<()> { Ok(()) } +/// The iterator of a poll with `items=`. It reads a page when it runs out of items, and keeps the +/// request that resumes after the last item that it yielded. +const PYTHON_POLL_ITERATOR: &str = r#"_PollRequestT = typing.TypeVar("_PollRequestT") +_PollItemT = typing.TypeVar("_PollItemT") + + +@typing.final +class PollIterator(typing.Generic[_PollRequestT, _PollItemT]): + """Yields each item that a poll reads, page by page. + + ``next_request`` resumes after the last item that the iterator yielded. It is ``None`` + after the last page. + """ + + def __init__( + self, + request: _PollRequestT, + read: collections.abc.Callable[ + [_PollRequestT], + collections.abc.Awaitable[ + tuple[_PollRequestT | None, collections.abc.Sequence[_PollItemT]] + ], + ], + resume: collections.abc.Callable[[_PollRequestT, _PollItemT], _PollRequestT] + | None, + ) -> None: + self.next_request: _PollRequestT | None = request + self._read: collections.abc.Callable[ + [_PollRequestT], + collections.abc.Awaitable[ + tuple[_PollRequestT | None, collections.abc.Sequence[_PollItemT]] + ], + ] = read + self._resume: ( + collections.abc.Callable[[_PollRequestT, _PollItemT], _PollRequestT] | None + ) = resume + self._page: list[_PollItemT] = [] + self._page_request: _PollRequestT = request + self._page_next: _PollRequestT | None = None + + def __aiter__(self) -> PollIterator[_PollRequestT, _PollItemT]: + return self + + async def __anext__(self) -> _PollItemT: + while not self._page: + if self.next_request is None: + raise StopAsyncIteration + self._page_request = self.next_request + self._page_next, items = await self._read(self._page_request) + self._page = list(items) + if not self._page: + self.next_request = self._page_next + item = self._page.pop(0) + # After a page's last item, the next read starts where the page helper says. + if not self._page: + self.next_request = self._page_next + elif self._resume is not None: + self.next_request = self._resume(self._page_request, item) + return item +"#; + /// The Nexus client that a client extension starts operations with. It has the /// `start_operation` of the Workflow Nexus client. const PYTHON_CLIENT_NEXUS_PROTOCOL: &str = r#"_OutputT = typing.TypeVar("_OutputT") @@ -8690,7 +8823,18 @@ impl PythonTarget { /// Starts the `def` line of an operation, with `self` for a method. fn render_operation_signature_start(output: &mut String, target: PythonTarget, name: &str) { - start_line(output, target, "async def "); + render_operation_signature_start_with(output, target, name, "async def "); +} + +/// Starts a function of the target with `keyword`, such as `def ` for a function that returns an +/// iterator at once. +fn render_operation_signature_start_with( + output: &mut String, + target: PythonTarget, + name: &str, + keyword: &str, +) { + start_line(output, target, keyword); output.push_str(name); output.push_str("(\n"); if target == PythonTarget::Client { @@ -8889,7 +9033,13 @@ fn poll_function_name(name: &str) -> String { format!("iter_{name}") } -fn poll_function_summary(name: &str) -> String { +fn poll_function_summary(name: &str, operation: &RenderedOperation<'_>) -> String { + if operation.poll_items.is_some() { + return format!( + "Yield each item that ``{name}`` reads, page by page. The iterator's \ + ``next_request`` resumes after the last item that it yielded." + ); + } format!( "Yield each response of ``{name}``. The poll helper gives the next request, and the \ iteration stops when it gives ``None``." @@ -8905,7 +9055,19 @@ fn render_operation_poll_function( next_expr: &str, ) { let body_indent = target.body_indent(); - render_operation_signature_start(output, target, &poll_function_name(&operation.attr_name)); + let page_annotation = match target { + PythonTarget::Workflow => { + temporalio_workflow_type_annotation(service, &operation.output_annotation) + } + PythonTarget::Client => operation.output_annotation.clone(), + }; + let (keyword, annotation) = operation.poll_iterator_signature(&page_annotation); + render_operation_signature_start_with( + output, + target, + &poll_function_name(&operation.attr_name), + keyword, + ); operation.render_public_runtime_type_parameters(output, &format!("{} ", target.indent())); let unpacked_input = operation.target_unpacked_input(target); if let Some(unpacked_input) = unpacked_input { @@ -8934,19 +9096,13 @@ fn render_operation_poll_function( let keyword_only = unpacked_input.is_some_and(|input| !input.parameters.is_empty()); render_rpc_option_parameters(output, " ", keyword_only); } - start_line(output, target, ") -> collections.abc.AsyncIterator["); - match target { - PythonTarget::Workflow => output.push_str(&temporalio_workflow_type_annotation( - service, - &operation.output_annotation, - )), - PythonTarget::Client => output.push_str(&operation.output_annotation), - } - output.push_str("]:\n"); + start_line(output, target, ") -> "); + output.push_str(&annotation); + output.push_str(":\n"); render_python_docstring( output, body_indent, - Some(&poll_function_summary(&operation.attr_name)), + Some(&poll_function_summary(&operation.attr_name, operation)), &[], None, false, @@ -8983,7 +9139,46 @@ fn render_operation_poll_function( ], PythonTarget::Client => vec![format!("response = await {call}")], }; - render_poll_loop(output, body_indent, &response, next_expr); + render_poll_body(output, body_indent, &response, next_expr, operation); +} + +/// The body of a poll iterator, after the first request. `response` is the code that gets the +/// response to `request`. An item poll hands a page reader to a `PollIterator`. +fn render_poll_body( + output: &mut String, + indent: &str, + response: &[String], + next_expr: &str, + operation: &RenderedOperation<'_>, +) { + let Some((items_attr, item_annotation)) = &operation.poll_items else { + render_poll_loop(output, indent, response, next_expr); + return; + }; + let request = operation.poll_request_annotation(); + output.push_str(&format!( + "\n{indent}async def read(\n{indent} request: {request},\n{indent}) -> " + )); + output.push_str(&format!( + "tuple[{request} | None, collections.abc.Sequence[{item_annotation}]]:\n" + )); + for line in response { + output.push_str(&format!("{indent} {line}\n")); + } + output.push_str(&format!( + "{indent} return {next_expr}(request, response), response.{items_attr} or []\n\n" + )); + output.push_str(&format!( + "{indent}return PollIterator(request, read, {})\n", + operation.poll_resume_expr.as_deref().unwrap_or("None") + )); +} + +/// Imports the package's `PollIterator`, when the body has an item poll. +fn render_poll_iterator_import(output: &mut String, body: &str, module: &str) { + if body.contains("PollIterator(") { + output.push_str(&format!("from {module} import PollIterator\n")); + } } /// The loop of a poll iterator, after the first request. `response` is the code that gets the diff --git a/src/parser/json_schema.rs b/src/parser/json_schema.rs index 73cebaa06..23bb614a2 100644 --- a/src/parser/json_schema.rs +++ b/src/parser/json_schema.rs @@ -6186,6 +6186,8 @@ fn build_operation( serialization_context: LanguageStringSpec::default(), poll_next: LanguageStringSpec::default(), poll_target: None, + poll_items: None, + poll_resume: LanguageStringSpec::default(), asynchronous: false, data: (), }) diff --git a/src/parser/wit.rs b/src/parser/wit.rs index 51af58de9..ea252b4d2 100644 --- a/src/parser/wit.rs +++ b/src/parser/wit.rs @@ -3518,6 +3518,8 @@ fn build_operation( )?; let serialization_context = build_operation_serialization_context(&directives, path, &context)?; let (poll_next, poll_target) = build_operation_poll(&directives, path, &context)?; + let (poll_items, poll_resume) = + build_operation_poll_items(resolve, function, &directives, path, &context)?; let output = if let Some(output_type) = function.result.as_ref() { Some( find_operation_type_spec(resolve, output_type, path, &context)?.ok_or_else(|| { @@ -3559,6 +3561,8 @@ fn build_operation( serialization_context, poll_next, poll_target, + poll_items, + poll_resume, asynchronous: matches!(function.kind, FunctionKind::AsyncFreestanding), data: (), }) @@ -3759,6 +3763,66 @@ fn build_operation_poll( Ok((spec, target)) } +/// The response field whose items a poll yields, by key, and the helper that resumes a read +/// after one item. +fn build_operation_poll_items( + resolve: &Resolve, + function: &Function, + directives: &[Directive], + path: &Path, + context: &str, +) -> Result<(Option, LanguageStringSpec)> { + let Some(directive) = directive(directives, "poll", path, context)? else { + return Ok((None, LanguageStringSpec::default())); + }; + let invalid = |reason: String| Error::InvalidWitDirective { + path: path.to_path_buf(), + context: context.to_string(), + directive: "@nexus.poll".to_string(), + reason, + }; + let resume = directive_language_string_for_key(directive, "resume"); + for helper in resume.default.iter().chain(resume.by_language.values()) { + if !is_valid_support_helper_path(helper) { + return Err(invalid(format!( + "invalid `resume` support helper `{helper}`" + ))); + } + } + let Some(items) = directive.value("items") else { + if resume.default.is_some() || !resume.by_language.is_empty() { + return Err(invalid("`resume` needs `items`".to_string())); + } + return Ok((None, resume)); + }; + let record = function + .result + .and_then(|ty| flattened_record(resolve, ty)) + .ok_or_else(|| { + invalid("`items` needs an operation whose result is a record".to_string()) + })?; + let field = record + .fields + .iter() + .find(|field| field.name == items) + .ok_or_else(|| invalid(format!("`{items}` isn't a field of the response")))?; + let mut ty = field.ty; + let is_list = loop { + let Type::Id(id) = ty else { + break false; + }; + match &resolve.types[id].kind { + TypeDefKind::Option(inner) | TypeDefKind::Type(inner) => ty = *inner, + TypeDefKind::List(_) => break true, + _ => break false, + } + }; + if !is_list { + return Err(invalid(format!("`{items}` isn't a list"))); + } + Ok((Some(record_field_key(field, path, context)?), resume)) +} + pub(crate) fn select_world( resolve: &Resolve, package_id: PackageId, diff --git a/src/planning/type_planning.rs b/src/planning/type_planning.rs index b03b40a64..219891d0a 100644 --- a/src/planning/type_planning.rs +++ b/src/planning/type_planning.rs @@ -706,6 +706,8 @@ impl<'a> TypePlanningContext<'a> { serialization_context: operation.serialization_context.clone(), poll_next: operation.poll_next.clone(), poll_target: operation.poll_target, + poll_items: operation.poll_items.clone(), + poll_resume: operation.poll_resume.clone(), asynchronous: operation.asynchronous, data: PlannedOperationData { output_resource_return: plan_operation_resource_return( @@ -1993,6 +1995,8 @@ interface alerts { serialization_context: LanguageStringSpec::default(), poll_next: LanguageStringSpec::default(), poll_target: None, + poll_items: None, + poll_resume: LanguageStringSpec::default(), asynchronous: false, data: (), }], diff --git a/src/spec.rs b/src/spec.rs index e5ce0fc23..2e63e458b 100644 --- a/src/spec.rs +++ b/src/spec.rs @@ -674,6 +674,12 @@ pub struct OperationSpec { pub poll_next: F::Text, /// The only target whose API gets the poll. Without it, every target gets the poll. pub poll_target: Option, + /// The key of the response's list field whose items the poll yields (`items=`). Without + /// it, the poll yields each response. + pub poll_items: Option, + /// The support helper that turns a page's request and one of its items into the request that + /// resumes after that item (`resume=`). + pub poll_resume: F::Text, /// Whether the operation can run for a long time, so a client can start it and wait later. pub asynchronous: bool, pub data: F::OperationData, @@ -714,6 +720,8 @@ impl OperationSpec { serialization_context: map.map_text(self.serialization_context), poll_next: map.map_text(self.poll_next), poll_target: self.poll_target, + poll_items: self.poll_items, + poll_resume: map.map_text(self.poll_resume), asynchronous: self.asynchronous, data, }