From 6a4528420eccf38b1c9bf80962cccda0632a6eea Mon Sep 17 00:00:00 2001 From: Thomas Grainger Date: Sun, 19 May 2024 20:45:00 +0100 Subject: [PATCH] avoid ResourceWarning in `DataBody.__aiter__` (#302) --- src/quart/wrappers/response.py | 20 +++++++++++++++----- 1 file changed, 15 insertions(+), 5 deletions(-) diff --git a/src/quart/wrappers/response.py b/src/quart/wrappers/response.py index 2fe27c2..065eabf 100644 --- a/src/quart/wrappers/response.py +++ b/src/quart/wrappers/response.py @@ -77,11 +77,8 @@ async def __aenter__(self) -> DataBody: async def __aexit__(self, exc_type: type, exc_value: BaseException, tb: TracebackType) -> None: pass - def __aiter__(self) -> AsyncIterator: - async def _aiter() -> AsyncGenerator[bytes, None]: - yield self.data[self.begin : self.end] - - return _aiter() + def __aiter__(self) -> AsyncIterator[bytes]: + return _DataBodyGen(self) async def make_conditional(self, begin: int, end: int | None) -> int: self.begin = begin @@ -91,6 +88,19 @@ async def make_conditional(self, begin: int, end: int | None) -> int: return len(self.data) +class _DataBodyGen(AsyncIterator[bytes]): + def __init__(self, data_body: DataBody): + self._data_body = data_body + self._iterated = False + + async def __anext__(self) -> bytes: + if self._iterated: + raise StopAsyncIteration + + self._iterated = True + return self._data_body.data[self._data_body.begin : self._data_body.end] + + class IterableBody(ResponseBody): def __init__(self, iterable: AsyncGenerator[bytes, None] | Iterable) -> None: self.iter: AsyncGenerator[bytes, None]