Skip to content

Commit 11b8b09

Browse files
yanyongyuyanyongyu
authored
✨ Feature: support streaming response data (#225)
1 parent fb4d115 commit 11b8b09

5 files changed

Lines changed: 89 additions & 5 deletions

File tree

‎codegen/templates/rest/_param.py.jinja‎

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -37,6 +37,7 @@
3737
{{ header_params(endpoint) }}
3838
{{ cookie_params(endpoint) }}
3939
headers: Optional[Mapping[str, str]] = None,
40+
stream: bool = False,
4041
{{ endpoint.request_body.get_raw_definition() }}
4142
{% endmacro %}
4243

@@ -48,6 +49,7 @@ headers: Optional[Mapping[str, str]] = None,
4849
{{ cookie_params(endpoint) }}
4950
data: UnsetType = UNSET,
5051
headers: Optional[Mapping[str, str]] = None,
52+
stream: bool = False,
5153
{{ body_params(model, endpoint.param_names) }}
5254
{% endmacro %}
5355

@@ -58,6 +60,7 @@ headers: Optional[Mapping[str, str]] = None,
5860
{{ header_params(endpoint) }}
5961
{{ cookie_params(endpoint) }}
6062
headers: Optional[Mapping[str, str]] = None,
63+
stream: bool = False,
6164
{%- if endpoint.request_body %}
6265
{{ endpoint.request_body.get_endpoint_definition() }},
6366
{%- if endpoint.request_body.allowed_models %}

‎codegen/templates/rest/_request.py.jinja‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -103,6 +103,7 @@ headers=exclude_unset(headers),
103103
{% if endpoint.cookie_params %}
104104
cookies=exclude_unset(cookies),
105105
{% endif %}
106+
stream=stream,
106107
{% if endpoint.success_response and endpoint.success_response.response_schema %}
107108
response_model={{ build_response_model(endpoint.success_response) }},
108109
{% endif %}

‎githubkit/core.py‎

Lines changed: 26 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -286,6 +286,7 @@ def _request(
286286
json: Optional[Any] = None,
287287
headers: Optional[HeaderTypes] = None,
288288
cookies: Optional[CookieTypes] = None,
289+
stream: bool = False,
289290
) -> httpx.Response:
290291
with self.get_sync_client() as client:
291292
request = client.build_request(
@@ -301,7 +302,7 @@ def _request(
301302
)
302303
with self.config.throttler.acquire(request):
303304
try:
304-
return client.send(request)
305+
return client.send(request, stream=stream)
305306
except httpx.TimeoutException as e:
306307
raise RequestTimeout(e) from e
307308
except Exception as e:
@@ -320,6 +321,7 @@ async def _arequest(
320321
json: Optional[Any] = None,
321322
headers: Optional[HeaderTypes] = None,
322323
cookies: Optional[CookieTypes] = None,
324+
stream: bool = False,
323325
) -> httpx.Response:
324326
async with (
325327
self.get_async_client() as client,
@@ -337,7 +339,7 @@ async def _arequest(
337339
)
338340
async with self.config.throttler.async_acquire(request):
339341
try:
340-
return await client.send(request)
342+
return await client.send(request, stream=stream)
341343
except httpx.TimeoutException as e:
342344
raise RequestTimeout(e) from e
343345
except Exception as e:
@@ -360,13 +362,17 @@ def _check(
360362
error_models: Optional[Mapping[str, type]] = None,
361363
) -> Response[Any]: ...
362364

365+
def _check_is_error(self, response: httpx.Response) -> bool:
366+
"""Check if the response is an error."""
367+
return response.is_error
368+
363369
def _check(
364370
self,
365371
response: httpx.Response,
366372
response_model: Union[type[T], UnsetType] = UNSET,
367373
error_models: Optional[Mapping[str, type]] = None,
368374
) -> Union[Response[T], Response[Any]]:
369-
if response.is_error:
375+
if self._check_is_error(response):
370376
error_models = error_models or {}
371377
status_code = str(response.status_code)
372378

@@ -386,7 +392,7 @@ def _check(
386392
if response.status_code in (403, 429):
387393
self._check_rate_limit(resp)
388394

389-
if response.is_error:
395+
if self._check_is_error(response):
390396
raise RequestFailed(resp)
391397
return resp
392398

@@ -453,6 +459,7 @@ def request(
453459
json: Optional[Any] = None,
454460
headers: Optional[HeaderTypes] = None,
455461
cookies: Optional[CookieTypes] = None,
462+
stream: bool = False,
456463
response_model: type[T],
457464
error_models: Optional[Mapping[str, type]] = None,
458465
) -> Response[T]: ...
@@ -470,6 +477,7 @@ def request(
470477
json: Optional[Any] = None,
471478
headers: Optional[HeaderTypes] = None,
472479
cookies: Optional[CookieTypes] = None,
480+
stream: bool = False,
473481
response_model: UnsetType = UNSET,
474482
error_models: Optional[Mapping[str, type]] = None,
475483
) -> Response[Any]: ...
@@ -486,6 +494,7 @@ def request(
486494
json: Optional[Any] = None,
487495
headers: Optional[HeaderTypes] = None,
488496
cookies: Optional[CookieTypes] = None,
497+
stream: bool = False,
489498
response_model: Union[type[T], UnsetType] = UNSET,
490499
error_models: Optional[Mapping[str, type]] = None,
491500
) -> Union[Response[T], Response[Any]]:
@@ -507,7 +516,12 @@ def request(
507516
json=json,
508517
headers=headers,
509518
cookies=cookies,
519+
stream=stream,
510520
)
521+
if self._check_is_error(raw_resp) and stream:
522+
# if the response is an error and stream is True,
523+
# we need to read the response first
524+
raw_resp.read()
511525
return self._check(raw_resp, response_model, error_models)
512526
except GitHubException as e:
513527
if self.config.auto_retry is None:
@@ -535,6 +549,7 @@ async def arequest(
535549
json: Optional[Any] = None,
536550
headers: Optional[HeaderTypes] = None,
537551
cookies: Optional[CookieTypes] = None,
552+
stream: bool = False,
538553
response_model: type[T],
539554
error_models: Optional[Mapping[str, type]] = None,
540555
) -> Response[T]: ...
@@ -552,6 +567,7 @@ async def arequest(
552567
json: Optional[Any] = None,
553568
headers: Optional[HeaderTypes] = None,
554569
cookies: Optional[CookieTypes] = None,
570+
stream: bool = False,
555571
response_model: UnsetType = UNSET,
556572
error_models: Optional[Mapping[str, type]] = None,
557573
) -> Response[Any]: ...
@@ -568,6 +584,7 @@ async def arequest(
568584
json: Optional[Any] = None,
569585
headers: Optional[HeaderTypes] = None,
570586
cookies: Optional[CookieTypes] = None,
587+
stream: bool = False,
571588
response_model: Union[type[T], UnsetType] = UNSET,
572589
error_models: Optional[Mapping[str, type]] = None,
573590
) -> Union[Response[T], Response[Any]]:
@@ -589,7 +606,12 @@ async def arequest(
589606
json=json,
590607
headers=headers,
591608
cookies=cookies,
609+
stream=stream,
592610
)
611+
if self._check_is_error(raw_resp) and stream:
612+
# if the response is an error and stream is True,
613+
# we need to read the response first
614+
await raw_resp.aread()
593615
return self._check(raw_resp, response_model, error_models)
594616
except GitHubException as e:
595617
if self.config.auto_retry is None:

‎githubkit/response.py‎

Lines changed: 32 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1,4 +1,5 @@
1-
from typing import Any, Generic
1+
from collections.abc import AsyncIterator, Iterator
2+
from typing import Any, Generic, Optional
23
from typing_extensions import TypeVar
34

45
import httpx
@@ -92,3 +93,33 @@ def json(self, **kwargs: Any) -> JT:
9293
@property
9394
def parsed_data(self) -> MT:
9495
return type_validate_json(self._data_model, self.content)
96+
97+
def iter_bytes(self, chunk_size: Optional[int] = None) -> Iterator[bytes]:
98+
yield from self._response.iter_bytes(chunk_size=chunk_size)
99+
100+
def iter_text(self, chunk_size: Optional[int] = None) -> Iterator[str]:
101+
yield from self._response.iter_text(chunk_size=chunk_size)
102+
103+
def iter_lines(self) -> Iterator[str]:
104+
yield from self._response.iter_lines()
105+
106+
def iter_raw(self, chunk_size: Optional[int] = None) -> Iterator[bytes]:
107+
yield from self._response.iter_raw(chunk_size=chunk_size)
108+
109+
async def aiter_bytes(
110+
self, chunk_size: Optional[int] = None
111+
) -> AsyncIterator[bytes]:
112+
async for chunk in self._response.aiter_bytes(chunk_size=chunk_size):
113+
yield chunk
114+
115+
async def aiter_text(self, chunk_size: Optional[int] = None) -> AsyncIterator[str]:
116+
async for chunk in self._response.aiter_text(chunk_size=chunk_size):
117+
yield chunk
118+
119+
async def aiter_lines(self) -> AsyncIterator[str]:
120+
async for line in self._response.aiter_lines():
121+
yield line
122+
123+
async def aiter_raw(self, chunk_size: Optional[int] = None) -> AsyncIterator[bytes]:
124+
async for chunk in self._response.aiter_raw(chunk_size=chunk_size):
125+
yield chunk

‎tests/test_rest/test_call.py‎

Lines changed: 27 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1,5 +1,6 @@
11
from functools import partial
22

3+
from httpx import ResponseNotRead
34
import pytest
45

56
from githubkit import GitHub
@@ -8,6 +9,7 @@
89

910
OWNER = "yanyongyu"
1011
REPO = "githubkit"
12+
REF = "master"
1113
ISSUE_COUNT_QUERY = """
1214
query($owner: String!, $repo: String!) {
1315
repository(owner: $owner, name: $repo) {
@@ -63,6 +65,31 @@ async def test_async_call_with_raw_body(g: GitHub):
6365
assert isinstance(resp.text, str)
6466

6567

68+
def test_call_streaming(g: GitHub):
69+
resp = g.rest.repos.download_tarball_archive(OWNER, REPO, REF, stream=True)
70+
71+
with pytest.raises(ResponseNotRead):
72+
resp.content
73+
74+
for chunk in resp.iter_bytes():
75+
assert isinstance(chunk, bytes)
76+
assert len(chunk) > 0
77+
78+
79+
@pytest.mark.anyio
80+
async def test_async_call_streaming(g: GitHub):
81+
resp = await g.rest.repos.async_download_tarball_archive(
82+
OWNER, REPO, REF, stream=True
83+
)
84+
85+
with pytest.raises(ResponseNotRead):
86+
resp.content
87+
88+
async for chunk in resp.aiter_bytes():
89+
assert isinstance(chunk, bytes)
90+
assert len(chunk) > 0
91+
92+
6693
def test_paginate(g: GitHub):
6794
paginator = g.rest.paginate(
6895
g.rest.issues.list_for_repo, owner=OWNER, repo=REPO, state="all", per_page=50

0 commit comments

Comments
 (0)