✨ close and map stream exception

yanyongyu · web-flow · commit 59e6135d0319 · 2025-06-23T12:14:11.000Z
diff --git a/githubkit/response.py b/githubkit/response.py
@@ -1,10 +1,12 @@
 from collections.abc import AsyncIterator, Iterator
+from contextlib import asynccontextmanager, contextmanager
 from typing import Any, Generic, Optional
 from typing_extensions import TypeVar
 
 import httpx
 
 from .compat import type_validate_json
+from .exception import RequestError, RequestTimeout
 
 MT = TypeVar("MT", default=Any)
 JT = TypeVar("JT", default=Any)
@@ -94,32 +96,62 @@ def json(self, **kwargs: Any) -> JT:
     def parsed_data(self) -> MT:
         return type_validate_json(self._data_model, self.content)
 
+    @contextmanager
+    def _catch_and_close(self):
+        try:
+            yield
+        except httpx.TimeoutException as e:
+            raise RequestTimeout(e) from e
+        except Exception as e:
+            raise RequestError(e) from e
+        finally:
+            self._response.close()
+
+    @asynccontextmanager
+    async def _acatch_and_close(self):
+        try:
+            yield
+        except httpx.TimeoutException as e:
+            raise RequestTimeout(e) from e
+        except Exception as e:
+            raise RequestError(e) from e
+        finally:
+            await self._response.aclose()
+
     def iter_bytes(self, chunk_size: Optional[int] = None) -> Iterator[bytes]:
-        yield from self._response.iter_bytes(chunk_size=chunk_size)
+        with self._catch_and_close():
+            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)
+        with self._catch_and_close():
+            yield from self._response.iter_text(chunk_size=chunk_size)
 
     def iter_lines(self) -> Iterator[str]:
-        yield from self._response.iter_lines()
+        with self._catch_and_close():
+            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)
+        with self._catch_and_close():
+            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 with self._acatch_and_close():
+            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 with self._acatch_and_close():
+            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 with self._acatch_and_close():
+            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
+        async with self._acatch_and_close():
+            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
@@ -1,6 +1,5 @@
 from functools import partial
 
-from httpx import ResponseNotRead
 import pytest
 
 from githubkit import GitHub
@@ -65,31 +64,6 @@ 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
