diff options
Diffstat (limited to 'discord/iterators.py')
| -rw-r--r-- | discord/iterators.py | 70 |
1 files changed, 44 insertions, 26 deletions
diff --git a/discord/iterators.py b/discord/iterators.py index 63a8776d..91470d80 100644 --- a/discord/iterators.py +++ b/discord/iterators.py @@ -27,23 +27,26 @@ DEALINGS IN THE SOFTWARE. import sys import asyncio import aiohttp +import datetime + +from .errors import NoMoreMessages +from .utils import time_snowflake from .message import Message from .object import Object PY35 = sys.version_info >= (3, 5) - class LogsFromIterator: - """Iterator for recieving logs. + """Iterator for receiving logs. - The messages endpoint has two behaviors we care about here: + The messages endpoint has two behaviours we care about here: If `before` is specified, the messages endpoint returns the `limit` newest messages before `before`, sorted with newest first. For filling over - 100 messages, update the `before` parameter to the oldest message recieved. + 100 messages, update the `before` parameter to the oldest message received. Messages will be returned in order by time. If `after` is specified, it returns the `limit` oldest messages after `after`, sorted with newest first. For filling over 100 messages, update the - `after` parameter to the newest message recieved. If messages are not + `after` parameter to the newest message received. If messages are not reversed, they will be out of order (99-0, 199-100, so on) A note that if both before and after are specified, before is ignored by the @@ -51,8 +54,7 @@ class LogsFromIterator: Parameters ----------- - client : class:`Client` - channel : class:`Channel` + channel: class:`Channel` Channel from which to request logs limit : int Maximum number of messages to retrieve @@ -63,24 +65,37 @@ class LogsFromIterator: around : :class:`Message` or id-like Message around which all messages must be. Limit max 101. Note that if limit is an even number, this will return at most limit+1 messages. - reverse : bool + reverse: bool If set to true, return messages in oldest->newest order. Recommended when using with "after" queries with limit over 100, otherwise messages - will be out of order. Defaults to False for backwards compatability. + will be out of order. """ - def __init__(self, client, channel, limit, - before=None, after=None, around=None, reverse=False): - self.client = client + def __init__(self, channel, limit, + before=None, after=None, around=None, reverse=None): + + if isinstance(before, datetime.datetime): + before = Object(id=time_snowflake(before, high=False)) + if isinstance(after, datetime.datetime): + after = Object(id=time_snowflake(after, high=True)) + if isinstance(around, datetime.datetime): + around = Object(id=time_snowflake(around)) + self.channel = channel + self.ctx = channel._state + self.logs_from = channel._state.http.logs_from self.limit = limit self.before = before self.after = after self.around = around - self.reverse = reverse + + if reverse is None: + self.reverse = after is not None + else: + self.reverse = reverse + self._filter = None # message dict -> bool self.messages = asyncio.Queue() - self.ctx = client.connection.ctx if self.around: if self.limit > 101: @@ -92,29 +107,32 @@ class LogsFromIterator: self._retrieve_messages = self._retrieve_messages_around_strategy if self.before and self.after: - self._filter = lambda m: self.after.id < m['id'] < self.before.id + self._filter = lambda m: self.after.id < int(m['id']) < self.before.id elif self.before: - self._filter = lambda m: m['id'] < self.before.id + self._filter = lambda m: int(m['id']) < self.before.id elif self.after: - self._filter = lambda m: self.after.id < m['id'] + self._filter = lambda m: self.after.id < int(m['id']) elif self.before and self.after: if self.reverse: self._retrieve_messages = self._retrieve_messages_after_strategy - self._filter = lambda m: m['id'] < self.before.id + self._filter = lambda m: int(m['id']) < self.before.id else: self._retrieve_messages = self._retrieve_messages_before_strategy - self._filter = lambda m: m['id'] > self.after.id + self._filter = lambda m: int(m['id']) > self.after.id elif self.after: self._retrieve_messages = self._retrieve_messages_after_strategy else: self._retrieve_messages = self._retrieve_messages_before_strategy @asyncio.coroutine - def iterate(self): + def get(self): if self.messages.empty(): yield from self.fill_messages() - return self.messages.get_nowait() + try: + return self.messages.get_nowait() + except asyncio.QueueEmpty: + raise NoMoreMessages() @asyncio.coroutine def fill_messages(self): @@ -136,7 +154,7 @@ class LogsFromIterator: @asyncio.coroutine def _retrieve_messages_before_strategy(self, retrieve): """Retrieve messages using before parameter.""" - data = yield from self.client._logs_from(self.channel, retrieve, before=self.before) + data = yield from self.logs_from(self.channel.id, retrieve, before=getattr(self.before, 'id', None)) if len(data): self.limit -= retrieve self.before = Object(id=int(data[-1]['id'])) @@ -145,7 +163,7 @@ class LogsFromIterator: @asyncio.coroutine def _retrieve_messages_after_strategy(self, retrieve): """Retrieve messages using after parameter.""" - data = yield from self.client._logs_from(self.channel, retrieve, after=self.after) + data = yield from self.logs_from(self.channel.id, retrieve, after=getattr(self.after, 'id', None)) if len(data): self.limit -= retrieve self.after = Object(id=int(data[0]['id'])) @@ -155,7 +173,7 @@ class LogsFromIterator: def _retrieve_messages_around_strategy(self, retrieve): """Retrieve messages using around parameter.""" if self.around: - data = yield from self.client._logs_from(self.channel, retrieve, around=self.around) + data = yield from self.logs_from(self.channel.id, retrieve, around=getattr(self.around, 'id', None)) self.around = None return data return [] @@ -168,9 +186,9 @@ class LogsFromIterator: @asyncio.coroutine def __anext__(self): try: - msg = yield from self.iterate() + msg = yield from self.get() return msg - except asyncio.QueueEmpty: + except NoMoreMessages: # if we're still empty at this point... # we didn't get any new messages so stop looping raise StopAsyncIteration() |