Implement Streamable

Author: patrick91

e2e/app.py

@@ -1,14 +1,26 @@
 import asyncio  # noqa: INP001
+import random
 
 import strawberry
 from strawberry.schema.config import StrawberryConfig
 
 
+@strawberry.type
+class Author:
+    id: strawberry.ID
+    name: str
+
+
 @strawberry.type
 class Comment:
     id: strawberry.ID
     content: str
 
+    @strawberry.field
+    async def author(self) -> Author:
+        await asyncio.sleep(random.randint(0, 2))  # noqa: S311
+        return Author(id=strawberry.ID("Author:1"), name="John Doe")
+
 
 @strawberry.type
 class BlogPost:
@@ -17,12 +29,11 @@ class BlogPost:
     content: str
 
     @strawberry.field
-    async def comments(self) -> list[Comment]:
-        await asyncio.sleep(2)
-        return [
-            Comment(id=strawberry.ID("Comment:1"), content="Great post!"),
-            Comment(id=strawberry.ID("Comment:2"), content="Thanks for sharing!"),
-        ]
+    async def comments(self) -> strawberry.Streamable[Comment]:
+        for x in range(5):
+            await asyncio.sleep(random.choice([0, 0.5, 1, 1.5, 2]))  # noqa: S311
+
+            yield Comment(id=strawberry.ID(f"Comment:{x}"), content="Great post!")
 
 
 @strawberry.type

strawberry/__init__.py

@@ -11,6 +11,7 @@
 from .scalars import ID
 from .schema import Schema
 from .schema_directive import schema_directive
+from .streamable import Streamable
 from .types.arguments import argument
 from .types.auto import auto
 from .types.cast import cast
@@ -34,6 +35,7 @@
     "Parent",
     "Private",
     "Schema",
+    "Streamable",
     "argument",
     "asdict",
     "auto",

strawberry/annotation.py

@@ -17,6 +17,7 @@
 )
 from typing_extensions import Self, get_args, get_origin
 
+from strawberry.streamable import StrawberryStreamable
 from strawberry.types.base import (
     StrawberryList,
     StrawberryObjectDefinition,
@@ -141,6 +142,8 @@ def _resolve(self) -> Union[StrawberryType, type]:
 
         if self._is_lazy_type(evaled_type):
             return evaled_type
+        if self._is_streamable(evaled_type, args):
+            return self.create_list(list[evaled_type])
         if self._is_list(evaled_type):
             return self.create_list(evaled_type)
 
@@ -310,6 +313,10 @@ def _is_list(cls, annotation: Any) -> bool:
             or is_list
         )
 
+    @classmethod
+    def _is_streamable(cls, annotation: Any, args: list[Any]) -> bool:
+        return any(isinstance(arg, StrawberryStreamable) for arg in args)
+
     @classmethod
     def _is_strawberry_type(cls, evaled_type: Any) -> bool:
         # Prevent import cycles

strawberry/http/async_base_view.py

@@ -367,6 +367,8 @@ async def stream() -> AsyncGenerator[str, None]:
 
                 yield self.encode_multipart_data(response, "-")
 
+                all_pending = result.initial_result.pending
+
                 async for value in result.subsequent_results:
                     response = {
                         "hasNext": value.has_next,
@@ -382,11 +384,13 @@ async def stream() -> AsyncGenerator[str, None]:
                     if value.incremental:
                         incremental = []
 
+                        all_pending.extend(value.pending)
+
                         for incremental_value in value.incremental:
                             pending_value = next(
                                 (
                                     v
-                                    for v in result.initial_result.pending
+                                    for v in all_pending
                                     if v.id == incremental_value.id
                                 ),
                                 None,
@@ -397,8 +401,6 @@ async def stream() -> AsyncGenerator[str, None]:
                             incremental.append(
                                 {
                                     **incremental_value.formatted,
-                                    # for Apollo
-                                    # content type is `multipart/mixed;deferSpec=20220824,application/json`
                                     "path": pending_value.path,
                                     "label": pending_value.label,
                                 }

strawberry/streamable.py

@@ -0,0 +1,36 @@
+from collections.abc import AsyncGenerator
+from typing import Annotated, TypeVar
+
+
+class StrawberryStreamable: ...
+
+
+T = TypeVar("T")
+
+Streamable = Annotated[AsyncGenerator[T, None], StrawberryStreamable()]
+"""Represents a list that can be streamed using @stream.
+
+Example:
+
+```python
+import strawberry
+from dataclasses import dataclass
+
+
+@strawberry.type
+class Comment:
+    id: strawberry.ID
+    content: str
+
+
+@strawberry.type
+class Article:
+    @strawberry.field
+    @staticmethod
+    async def comments() -> strawberry.Streamable[Comment]:
+        for comment in fetch_comments():
+            yield comment
+```
+"""
+
+__all__ = ["Streamable"]