peewee-async.manager.execute(Model.select()) loads everything into memory. Avoid that with this iterator

iter_select.py

"""
peewee-async.manager.execute(Model.select()) loads everything into memory. Avoid that with this iterator

Usage::
          
  q = Item.select()
  async for item in IterSelect(manager_instance, q):
    print(item.id)

DISCLAIMERS: Untested and seems slow
"""


from collections import AsyncIterable

import peewee
import peewee_async


class IterSelect(AsyncIterable):
    """
    Standard AsyncQueryWrapper loads all retrieved stuff into memory.
    This one will fetch results one by one.
    """

    def __init__(self, manager: peewee_async.Manager, query: peewee.SelectQuery):
        assert isinstance(query, peewee.SelectQuery), \
            ("Error, trying to run select coroutine with wrong query class %s" % str(query))
        self._query = manager._swap_database(query)
        self._cursor = None
        self._initialized = False
        self._result_wrapper = peewee_async.AsyncQueryWrapper._get_result_wrapper(self._query)

    async def __aiter__(self):
        return self

    async def __anext__(self):
        # TODO: Check situations when cursor is not released
        if not self._initialized and self._cursor is None:
            self._cursor = await peewee_async._execute_query_async(self._query)

        row = await self._cursor.fetchone()

        if not row:
            await self._cursor.release
            self._cursor = None
            self._result_wrapper = None
            raise StopAsyncIteration
        elif not self._initialized:
            self._result_wrapper.initialize(self._cursor.description)
            self._initialized = True

        obj = self._result_wrapper.process_row(row)
        return obj