Skip to content

Commit

Permalink
avoid ResourceWarning in DataBody.__aiter__ (#302)
Browse files Browse the repository at this point in the history
  • Loading branch information
graingert committed May 19, 2024
1 parent 425f685 commit 6a45284
Showing 1 changed file with 15 additions and 5 deletions.
20 changes: 15 additions & 5 deletions src/quart/wrappers/response.py
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand All @@ -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]
Expand Down

0 comments on commit 6a45284

Please sign in to comment.