diff --git a/packages/code-storage-python/pierre_storage/repo.py b/packages/code-storage-python/pierre_storage/repo.py index 945510f..f49248a 100644 --- a/packages/code-storage-python/pierre_storage/repo.py +++ b/packages/code-storage-python/pierre_storage/repo.py @@ -1,5 +1,6 @@ """Repository implementation for Pierre Git Storage SDK.""" +import contextlib import warnings from datetime import datetime, timezone from types import TracebackType @@ -51,15 +52,18 @@ class StreamingResponse: """Stream wrapper that keeps the HTTP client alive until closed.""" - def __init__(self, response: httpx.Response, client: httpx.AsyncClient) -> None: + def __init__(self, response: httpx.Response, client: httpx.AsyncClient, stream_context: Optional[contextlib.AbstractAsyncContextManager] = None) -> None: self._response = response self._client = client + self._stream_context = stream_context def __getattr__(self, name: str) -> Any: return getattr(self._response, name) async def aclose(self) -> None: await self._response.aclose() + if self._stream_context is not None: + await self._stream_context.__aexit__(None, None, None) await self._client.aclose() async def __aenter__(self) -> "StreamingResponse": @@ -264,7 +268,7 @@ async def get_file_stream( await client.aclose() raise - return StreamingResponse(response, client) + return StreamingResponse(response, client, stream_context) async def get_archive_stream( self, @@ -324,7 +328,7 @@ async def get_archive_stream( await client.aclose() raise - return StreamingResponse(response, client) + return StreamingResponse(response, client, stream_context) async def list_files( self, diff --git a/packages/code-storage-python/pyproject.toml b/packages/code-storage-python/pyproject.toml index 228736a..bd1640e 100644 --- a/packages/code-storage-python/pyproject.toml +++ b/packages/code-storage-python/pyproject.toml @@ -4,7 +4,7 @@ build-backend = "setuptools.build_meta" [project] name = "pierre-storage" -version = "1.4.3" +version = "1.4.4" description = "Pierre Git Storage SDK for Python" readme = "README.md" license = "MIT"