mirror of
https://github.com/LonamiWebs/Telethon.git
synced 2025-08-08 04:52:30 +00:00
Create a new RequestIter ABC to deal with iter methods
This should make it easier to maintain these methods, increase reusability, and get rid of the async_generator dependency. In the future, people could use this to more easily deal with raw API themselves.
This commit is contained in:
99
telethon/requestiter.py
Normal file
99
telethon/requestiter.py
Normal file
@@ -0,0 +1,99 @@
|
||||
import abc
|
||||
import asyncio
|
||||
import time
|
||||
|
||||
|
||||
class RequestIter(abc.ABC):
|
||||
"""
|
||||
Helper class to deal with requests that need offsets to iterate.
|
||||
|
||||
It has some facilities, such as automatically sleeping a desired
|
||||
amount of time between requests if needed (but not more).
|
||||
|
||||
Can be used synchronously if the event loop is not running and
|
||||
as an asynchronous iterator otherwise.
|
||||
|
||||
`limit` is the total amount of items that the iterator should return.
|
||||
This is handled on this base class, and will be always ``>= 0``.
|
||||
|
||||
`left` will be reset every time the iterator is used and will indicate
|
||||
the amount of items that should be emitted left, so that subclasses can
|
||||
be more efficient and fetch only as many items as they need.
|
||||
|
||||
Iterators may be used with ``reversed``, and their `reverse` flag will
|
||||
be set to ``True`` if that's the case. Note that if this flag is set,
|
||||
`buffer` should be filled in reverse too.
|
||||
"""
|
||||
def __init__(self, client, limit, *, reverse=False, wait_time=None, **kwargs):
|
||||
self.client = client
|
||||
self.reverse = reverse
|
||||
self.wait_time = wait_time
|
||||
self.kwargs = kwargs
|
||||
self.limit = max(float('inf') if limit is None else limit, 0)
|
||||
self.left = None
|
||||
self.buffer = None
|
||||
self.index = None
|
||||
self.total = None
|
||||
self.last_load = None
|
||||
|
||||
async def _init(self, **kwargs):
|
||||
"""
|
||||
Called when asynchronous initialization is necessary. All keyword
|
||||
arguments passed to `__init__` will be forwarded here, and it's
|
||||
preferable to use named arguments in the subclasses without defaults
|
||||
to avoid forgetting or misspelling any of them.
|
||||
|
||||
This method may ``raise StopAsyncIteration`` if it cannot continue.
|
||||
"""
|
||||
|
||||
async def __anext__(self):
|
||||
if self.buffer is ():
|
||||
await self._init(**self.kwargs)
|
||||
|
||||
if self.index == len(self.buffer):
|
||||
# asyncio will handle times <= 0 to sleep 0 seconds
|
||||
if self.wait_time:
|
||||
await asyncio.sleep(
|
||||
self.wait_time - (time.time() - self.last_load),
|
||||
loop=self.client.loop
|
||||
)
|
||||
self.last_load = time.time()
|
||||
|
||||
self.index = 0
|
||||
self.buffer = await self._load_next_chunk()
|
||||
|
||||
if not self.buffer:
|
||||
raise StopAsyncIteration
|
||||
|
||||
result = self.buffer[self.index]
|
||||
self.left -= 1
|
||||
self.index += 1
|
||||
return result
|
||||
|
||||
def __aiter__(self):
|
||||
self.buffer = ()
|
||||
self.index = 0
|
||||
self.last_load = 0
|
||||
self.left = self.limit
|
||||
return self
|
||||
|
||||
def __iter__(self):
|
||||
if self.client.loop.is_running():
|
||||
raise RuntimeError(
|
||||
'You must use "async for" if the event loop '
|
||||
'is running (i.e. you are inside an "async def")'
|
||||
)
|
||||
|
||||
raise NotImplementedError('lol!')
|
||||
|
||||
@abc.abstractmethod
|
||||
async def _load_next_chunk(self):
|
||||
"""
|
||||
Called when the next chunk is necessary.
|
||||
It should *always* return a `list`.
|
||||
"""
|
||||
raise NotImplementedError
|
||||
|
||||
def __reversed__(self):
|
||||
self.reverse = not self.reverse
|
||||
return self # __aiter__ will be called after, too
|
Reference in New Issue
Block a user