diff --git a/codegen/templates/rest/_param.py.jinja b/codegen/templates/rest/_param.py.jinja index a4f922aa74..91ba080756 100644 --- a/codegen/templates/rest/_param.py.jinja +++ b/codegen/templates/rest/_param.py.jinja @@ -37,6 +37,7 @@ {{ header_params(endpoint) }} {{ cookie_params(endpoint) }} headers: Optional[Mapping[str, str]] = None, +stream: bool = False, {{ endpoint.request_body.get_raw_definition() }} {% endmacro %} @@ -48,6 +49,7 @@ headers: Optional[Mapping[str, str]] = None, {{ cookie_params(endpoint) }} data: UnsetType = UNSET, headers: Optional[Mapping[str, str]] = None, +stream: bool = False, {{ body_params(model, endpoint.param_names) }} {% endmacro %} @@ -58,6 +60,7 @@ headers: Optional[Mapping[str, str]] = None, {{ header_params(endpoint) }} {{ cookie_params(endpoint) }} headers: Optional[Mapping[str, str]] = None, +stream: bool = False, {%- if endpoint.request_body %} {{ endpoint.request_body.get_endpoint_definition() }}, {%- if endpoint.request_body.allowed_models %} diff --git a/codegen/templates/rest/_request.py.jinja b/codegen/templates/rest/_request.py.jinja index e09adc62c6..fe6025efff 100644 --- a/codegen/templates/rest/_request.py.jinja +++ b/codegen/templates/rest/_request.py.jinja @@ -103,6 +103,7 @@ headers=exclude_unset(headers), {% if endpoint.cookie_params %} cookies=exclude_unset(cookies), {% endif %} +stream=stream, {% if endpoint.success_response and endpoint.success_response.response_schema %} response_model={{ build_response_model(endpoint.success_response) }}, {% endif %} diff --git a/githubkit/core.py b/githubkit/core.py index 78b21b5aed..40fac4d14e 100644 --- a/githubkit/core.py +++ b/githubkit/core.py @@ -286,6 +286,7 @@ def _request( json: Optional[Any] = None, headers: Optional[HeaderTypes] = None, cookies: Optional[CookieTypes] = None, + stream: bool = False, ) -> httpx.Response: with self.get_sync_client() as client: request = client.build_request( @@ -301,7 +302,7 @@ def _request( ) with self.config.throttler.acquire(request): try: - return client.send(request) + return client.send(request, stream=stream) except httpx.TimeoutException as e: raise RequestTimeout(e) from e except Exception as e: @@ -320,6 +321,7 @@ async def _arequest( json: Optional[Any] = None, headers: Optional[HeaderTypes] = None, cookies: Optional[CookieTypes] = None, + stream: bool = False, ) -> httpx.Response: async with ( self.get_async_client() as client, @@ -337,7 +339,7 @@ async def _arequest( ) async with self.config.throttler.async_acquire(request): try: - return await client.send(request) + return await client.send(request, stream=stream) except httpx.TimeoutException as e: raise RequestTimeout(e) from e except Exception as e: @@ -360,13 +362,17 @@ def _check( error_models: Optional[Mapping[str, type]] = None, ) -> Response[Any]: ... + def _check_is_error(self, response: httpx.Response) -> bool: + """Check if the response is an error.""" + return response.is_error + def _check( self, response: httpx.Response, response_model: Union[type[T], UnsetType] = UNSET, error_models: Optional[Mapping[str, type]] = None, ) -> Union[Response[T], Response[Any]]: - if response.is_error: + if self._check_is_error(response): error_models = error_models or {} status_code = str(response.status_code) @@ -386,7 +392,7 @@ def _check( if response.status_code in (403, 429): self._check_rate_limit(resp) - if response.is_error: + if self._check_is_error(response): raise RequestFailed(resp) return resp @@ -453,6 +459,7 @@ def request( json: Optional[Any] = None, headers: Optional[HeaderTypes] = None, cookies: Optional[CookieTypes] = None, + stream: bool = False, response_model: type[T], error_models: Optional[Mapping[str, type]] = None, ) -> Response[T]: ... @@ -470,6 +477,7 @@ def request( json: Optional[Any] = None, headers: Optional[HeaderTypes] = None, cookies: Optional[CookieTypes] = None, + stream: bool = False, response_model: UnsetType = UNSET, error_models: Optional[Mapping[str, type]] = None, ) -> Response[Any]: ... @@ -486,6 +494,7 @@ def request( json: Optional[Any] = None, headers: Optional[HeaderTypes] = None, cookies: Optional[CookieTypes] = None, + stream: bool = False, response_model: Union[type[T], UnsetType] = UNSET, error_models: Optional[Mapping[str, type]] = None, ) -> Union[Response[T], Response[Any]]: @@ -507,7 +516,12 @@ def request( json=json, headers=headers, cookies=cookies, + stream=stream, ) + if self._check_is_error(raw_resp) and stream: + # if the response is an error and stream is True, + # we need to read the response first + raw_resp.read() return self._check(raw_resp, response_model, error_models) except GitHubException as e: if self.config.auto_retry is None: @@ -535,6 +549,7 @@ async def arequest( json: Optional[Any] = None, headers: Optional[HeaderTypes] = None, cookies: Optional[CookieTypes] = None, + stream: bool = False, response_model: type[T], error_models: Optional[Mapping[str, type]] = None, ) -> Response[T]: ... @@ -552,6 +567,7 @@ async def arequest( json: Optional[Any] = None, headers: Optional[HeaderTypes] = None, cookies: Optional[CookieTypes] = None, + stream: bool = False, response_model: UnsetType = UNSET, error_models: Optional[Mapping[str, type]] = None, ) -> Response[Any]: ... @@ -568,6 +584,7 @@ async def arequest( json: Optional[Any] = None, headers: Optional[HeaderTypes] = None, cookies: Optional[CookieTypes] = None, + stream: bool = False, response_model: Union[type[T], UnsetType] = UNSET, error_models: Optional[Mapping[str, type]] = None, ) -> Union[Response[T], Response[Any]]: @@ -589,7 +606,12 @@ async def arequest( json=json, headers=headers, cookies=cookies, + stream=stream, ) + if self._check_is_error(raw_resp) and stream: + # if the response is an error and stream is True, + # we need to read the response first + await raw_resp.aread() return self._check(raw_resp, response_model, error_models) except GitHubException as e: if self.config.auto_retry is None: diff --git a/githubkit/response.py b/githubkit/response.py index a30cb9214d..6ae55c8860 100644 --- a/githubkit/response.py +++ b/githubkit/response.py @@ -1,4 +1,5 @@ -from typing import Any, Generic +from collections.abc import AsyncIterator, Iterator +from typing import Any, Generic, Optional from typing_extensions import TypeVar import httpx @@ -92,3 +93,33 @@ def json(self, **kwargs: Any) -> JT: @property def parsed_data(self) -> MT: return type_validate_json(self._data_model, self.content) + + def iter_bytes(self, chunk_size: Optional[int] = None) -> Iterator[bytes]: + yield from self._response.iter_bytes(chunk_size=chunk_size) + + def iter_text(self, chunk_size: Optional[int] = None) -> Iterator[str]: + yield from self._response.iter_text(chunk_size=chunk_size) + + def iter_lines(self) -> Iterator[str]: + yield from self._response.iter_lines() + + def iter_raw(self, chunk_size: Optional[int] = None) -> Iterator[bytes]: + yield from self._response.iter_raw(chunk_size=chunk_size) + + async def aiter_bytes( + self, chunk_size: Optional[int] = None + ) -> AsyncIterator[bytes]: + async for chunk in self._response.aiter_bytes(chunk_size=chunk_size): + yield chunk + + async def aiter_text(self, chunk_size: Optional[int] = None) -> AsyncIterator[str]: + async for chunk in self._response.aiter_text(chunk_size=chunk_size): + yield chunk + + async def aiter_lines(self) -> AsyncIterator[str]: + async for line in self._response.aiter_lines(): + yield line + + async def aiter_raw(self, chunk_size: Optional[int] = None) -> AsyncIterator[bytes]: + async for chunk in self._response.aiter_raw(chunk_size=chunk_size): + yield chunk diff --git a/tests/test_rest/test_call.py b/tests/test_rest/test_call.py index 177240160c..dd0f5bed54 100644 --- a/tests/test_rest/test_call.py +++ b/tests/test_rest/test_call.py @@ -1,5 +1,6 @@ from functools import partial +from httpx import ResponseNotRead import pytest from githubkit import GitHub @@ -8,6 +9,7 @@ OWNER = "yanyongyu" REPO = "githubkit" +REF = "master" ISSUE_COUNT_QUERY = """ query($owner: String!, $repo: String!) { repository(owner: $owner, name: $repo) { @@ -63,6 +65,31 @@ async def test_async_call_with_raw_body(g: GitHub): assert isinstance(resp.text, str) +def test_call_streaming(g: GitHub): + resp = g.rest.repos.download_tarball_archive(OWNER, REPO, REF, stream=True) + + with pytest.raises(ResponseNotRead): + resp.content + + for chunk in resp.iter_bytes(): + assert isinstance(chunk, bytes) + assert len(chunk) > 0 + + +@pytest.mark.anyio +async def test_async_call_streaming(g: GitHub): + resp = await g.rest.repos.async_download_tarball_archive( + OWNER, REPO, REF, stream=True + ) + + with pytest.raises(ResponseNotRead): + resp.content + + async for chunk in resp.aiter_bytes(): + assert isinstance(chunk, bytes) + assert len(chunk) > 0 + + def test_paginate(g: GitHub): paginator = g.rest.paginate( g.rest.issues.list_for_repo, owner=OWNER, repo=REPO, state="all", per_page=50