Kaylynn Morgan
3 years ago
committed by
GitHub
GitHub
8 changed files with 386 additions and 930 deletions
@ -1,753 +0,0 @@ |
|||||
""" |
|
||||
The MIT License (MIT) |
|
||||
|
|
||||
Copyright (c) 2015-present Rapptz |
|
||||
|
|
||||
Permission is hereby granted, free of charge, to any person obtaining a |
|
||||
copy of this software and associated documentation files (the "Software"), |
|
||||
to deal in the Software without restriction, including without limitation |
|
||||
the rights to use, copy, modify, merge, publish, distribute, sublicense, |
|
||||
and/or sell copies of the Software, and to permit persons to whom the |
|
||||
Software is furnished to do so, subject to the following conditions: |
|
||||
|
|
||||
The above copyright notice and this permission notice shall be included in |
|
||||
all copies or substantial portions of the Software. |
|
||||
|
|
||||
THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS |
|
||||
OR IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, |
|
||||
FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE |
|
||||
AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER |
|
||||
LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING |
|
||||
FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER |
|
||||
DEALINGS IN THE SOFTWARE. |
|
||||
""" |
|
||||
|
|
||||
from __future__ import annotations |
|
||||
|
|
||||
import asyncio |
|
||||
import datetime |
|
||||
from typing import Awaitable, TYPE_CHECKING, TypeVar, Optional, Any, Callable, Union, List, AsyncIterator |
|
||||
|
|
||||
from .errors import NoMoreItems |
|
||||
from .utils import snowflake_time, time_snowflake, maybe_coroutine |
|
||||
from .object import Object |
|
||||
from .audit_logs import AuditLogEntry |
|
||||
|
|
||||
__all__ = ( |
|
||||
'ReactionIterator', |
|
||||
'HistoryIterator', |
|
||||
'AuditLogIterator', |
|
||||
'GuildIterator', |
|
||||
'MemberIterator', |
|
||||
) |
|
||||
|
|
||||
if TYPE_CHECKING: |
|
||||
from .types.audit_log import ( |
|
||||
AuditLog as AuditLogPayload, |
|
||||
) |
|
||||
from .types.guild import ( |
|
||||
Guild as GuildPayload, |
|
||||
) |
|
||||
from .types.message import ( |
|
||||
Message as MessagePayload, |
|
||||
) |
|
||||
from .types.user import ( |
|
||||
PartialUser as PartialUserPayload, |
|
||||
) |
|
||||
|
|
||||
from .types.threads import ( |
|
||||
Thread as ThreadPayload, |
|
||||
) |
|
||||
|
|
||||
from .member import Member |
|
||||
from .user import User |
|
||||
from .message import Message |
|
||||
from .audit_logs import AuditLogEntry |
|
||||
from .guild import Guild |
|
||||
from .threads import Thread |
|
||||
from .abc import Snowflake |
|
||||
|
|
||||
T = TypeVar('T') |
|
||||
OT = TypeVar('OT') |
|
||||
_Func = Callable[[T], Union[OT, Awaitable[OT]]] |
|
||||
|
|
||||
OLDEST_OBJECT = Object(id=0) |
|
||||
|
|
||||
|
|
||||
class _AsyncIterator(AsyncIterator[T]): |
|
||||
__slots__ = () |
|
||||
|
|
||||
async def next(self) -> T: |
|
||||
raise NotImplementedError |
|
||||
|
|
||||
def get(self, **attrs: Any) -> Awaitable[Optional[T]]: |
|
||||
def predicate(elem: T): |
|
||||
for attr, val in attrs.items(): |
|
||||
nested = attr.split('__') |
|
||||
obj = elem |
|
||||
for attribute in nested: |
|
||||
obj = getattr(obj, attribute) |
|
||||
|
|
||||
if obj != val: |
|
||||
return False |
|
||||
return True |
|
||||
|
|
||||
return self.find(predicate) |
|
||||
|
|
||||
async def find(self, predicate: _Func[T, bool]) -> Optional[T]: |
|
||||
while True: |
|
||||
try: |
|
||||
elem = await self.next() |
|
||||
except NoMoreItems: |
|
||||
return None |
|
||||
|
|
||||
ret = await maybe_coroutine(predicate, elem) |
|
||||
if ret: |
|
||||
return elem |
|
||||
|
|
||||
def chunk(self, max_size: int) -> _ChunkedAsyncIterator[T]: |
|
||||
if max_size <= 0: |
|
||||
raise ValueError('async iterator chunk sizes must be greater than 0.') |
|
||||
return _ChunkedAsyncIterator(self, max_size) |
|
||||
|
|
||||
def map(self, func: _Func[T, OT]) -> _MappedAsyncIterator[OT]: |
|
||||
return _MappedAsyncIterator(self, func) |
|
||||
|
|
||||
def filter(self, predicate: _Func[T, bool]) -> _FilteredAsyncIterator[T]: |
|
||||
return _FilteredAsyncIterator(self, predicate) |
|
||||
|
|
||||
async def flatten(self) -> List[T]: |
|
||||
return [element async for element in self] |
|
||||
|
|
||||
async def __anext__(self) -> T: |
|
||||
try: |
|
||||
return await self.next() |
|
||||
except NoMoreItems: |
|
||||
raise StopAsyncIteration() |
|
||||
|
|
||||
|
|
||||
def _identity(x): |
|
||||
return x |
|
||||
|
|
||||
|
|
||||
class _ChunkedAsyncIterator(_AsyncIterator[List[T]]): |
|
||||
def __init__(self, iterator, max_size): |
|
||||
self.iterator = iterator |
|
||||
self.max_size = max_size |
|
||||
|
|
||||
async def next(self) -> List[T]: |
|
||||
ret: List[T] = [] |
|
||||
n = 0 |
|
||||
while n < self.max_size: |
|
||||
try: |
|
||||
item = await self.iterator.next() |
|
||||
except NoMoreItems: |
|
||||
if ret: |
|
||||
return ret |
|
||||
raise |
|
||||
else: |
|
||||
ret.append(item) |
|
||||
n += 1 |
|
||||
return ret |
|
||||
|
|
||||
|
|
||||
class _MappedAsyncIterator(_AsyncIterator[T]): |
|
||||
def __init__(self, iterator, func): |
|
||||
self.iterator = iterator |
|
||||
self.func = func |
|
||||
|
|
||||
async def next(self) -> T: |
|
||||
# this raises NoMoreItems and will propagate appropriately |
|
||||
item = await self.iterator.next() |
|
||||
return await maybe_coroutine(self.func, item) |
|
||||
|
|
||||
|
|
||||
class _FilteredAsyncIterator(_AsyncIterator[T]): |
|
||||
def __init__(self, iterator, predicate): |
|
||||
self.iterator = iterator |
|
||||
|
|
||||
if predicate is None: |
|
||||
predicate = _identity |
|
||||
|
|
||||
self.predicate = predicate |
|
||||
|
|
||||
async def next(self) -> T: |
|
||||
getter = self.iterator.next |
|
||||
pred = self.predicate |
|
||||
while True: |
|
||||
# propagate NoMoreItems similar to _MappedAsyncIterator |
|
||||
item = await getter() |
|
||||
ret = await maybe_coroutine(pred, item) |
|
||||
if ret: |
|
||||
return item |
|
||||
|
|
||||
|
|
||||
class ReactionIterator(_AsyncIterator[Union['User', 'Member']]): |
|
||||
def __init__(self, message, emoji, limit=100, after=None): |
|
||||
self.message = message |
|
||||
self.limit = limit |
|
||||
self.after = after |
|
||||
state = message._state |
|
||||
self.getter = state.http.get_reaction_users |
|
||||
self.state = state |
|
||||
self.emoji = emoji |
|
||||
self.guild = message.guild |
|
||||
self.channel_id = message.channel.id |
|
||||
self.users = asyncio.Queue() |
|
||||
|
|
||||
async def next(self) -> Union[User, Member]: |
|
||||
if self.users.empty(): |
|
||||
await self.fill_users() |
|
||||
|
|
||||
try: |
|
||||
return self.users.get_nowait() |
|
||||
except asyncio.QueueEmpty: |
|
||||
raise NoMoreItems() |
|
||||
|
|
||||
async def fill_users(self): |
|
||||
# this is a hack because >circular imports< |
|
||||
from .user import User |
|
||||
|
|
||||
if self.limit > 0: |
|
||||
retrieve = self.limit if self.limit <= 100 else 100 |
|
||||
|
|
||||
after = self.after.id if self.after else None |
|
||||
data: List[PartialUserPayload] = await self.getter( |
|
||||
self.channel_id, self.message.id, self.emoji, retrieve, after=after |
|
||||
) |
|
||||
|
|
||||
if data: |
|
||||
self.limit -= retrieve |
|
||||
self.after = Object(id=int(data[-1]['id'])) |
|
||||
|
|
||||
if self.guild is None or isinstance(self.guild, Object): |
|
||||
for element in reversed(data): |
|
||||
await self.users.put(User(state=self.state, data=element)) |
|
||||
else: |
|
||||
for element in reversed(data): |
|
||||
member_id = int(element['id']) |
|
||||
member = self.guild.get_member(member_id) |
|
||||
if member is not None: |
|
||||
await self.users.put(member) |
|
||||
else: |
|
||||
await self.users.put(User(state=self.state, data=element)) |
|
||||
|
|
||||
|
|
||||
class HistoryIterator(_AsyncIterator['Message']): |
|
||||
"""Iterator for receiving a channel's message history. |
|
||||
|
|
||||
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 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 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 |
|
||||
messages endpoint. |
|
||||
|
|
||||
Parameters |
|
||||
----------- |
|
||||
messageable: :class:`abc.Messageable` |
|
||||
Messageable class to retrieve message history from. |
|
||||
limit: :class:`int` |
|
||||
Maximum number of messages to retrieve |
|
||||
before: Optional[Union[:class:`abc.Snowflake`, :class:`datetime.datetime`]] |
|
||||
Message before which all messages must be. |
|
||||
after: Optional[Union[:class:`abc.Snowflake`, :class:`datetime.datetime`]] |
|
||||
Message after which all messages must be. |
|
||||
around: Optional[Union[:class:`abc.Snowflake`, :class:`datetime.datetime`]] |
|
||||
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. |
|
||||
oldest_first: Optional[:class:`bool`] |
|
||||
If set to ``True``, return messages in oldest->newest order. Defaults to |
|
||||
``True`` if `after` is specified, otherwise ``False``. |
|
||||
""" |
|
||||
|
|
||||
def __init__(self, messageable, limit, before=None, after=None, around=None, oldest_first=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)) |
|
||||
|
|
||||
if oldest_first is None: |
|
||||
self.reverse = after is not None |
|
||||
else: |
|
||||
self.reverse = oldest_first |
|
||||
|
|
||||
self.messageable = messageable |
|
||||
self.limit = limit |
|
||||
self.before = before |
|
||||
self.after = after or OLDEST_OBJECT |
|
||||
self.around = around |
|
||||
|
|
||||
self._filter = None # message dict -> bool |
|
||||
|
|
||||
self.state = self.messageable._state |
|
||||
self.logs_from = self.state.http.logs_from |
|
||||
self.messages = asyncio.Queue() |
|
||||
|
|
||||
if self.around: |
|
||||
if self.limit is None: |
|
||||
raise ValueError('history does not support around with limit=None') |
|
||||
if self.limit > 101: |
|
||||
raise ValueError("history max limit 101 when specifying around parameter") |
|
||||
elif self.limit == 101: |
|
||||
self.limit = 100 # Thanks discord |
|
||||
|
|
||||
self._retrieve_messages = self._retrieve_messages_around_strategy # type: ignore |
|
||||
if self.before and self.after: |
|
||||
self._filter = lambda m: self.after.id < int(m['id']) < self.before.id |
|
||||
elif self.before: |
|
||||
self._filter = lambda m: int(m['id']) < self.before.id |
|
||||
elif self.after: |
|
||||
self._filter = lambda m: self.after.id < int(m['id']) |
|
||||
else: |
|
||||
if self.reverse: |
|
||||
self._retrieve_messages = self._retrieve_messages_after_strategy # type: ignore |
|
||||
if self.before: |
|
||||
self._filter = lambda m: int(m['id']) < self.before.id |
|
||||
else: |
|
||||
self._retrieve_messages = self._retrieve_messages_before_strategy # type: ignore |
|
||||
if self.after and self.after != OLDEST_OBJECT: |
|
||||
self._filter = lambda m: int(m['id']) > self.after.id |
|
||||
|
|
||||
async def next(self) -> Message: |
|
||||
if self.messages.empty(): |
|
||||
await self.fill_messages() |
|
||||
|
|
||||
try: |
|
||||
return self.messages.get_nowait() |
|
||||
except asyncio.QueueEmpty: |
|
||||
raise NoMoreItems() |
|
||||
|
|
||||
def _get_retrieve(self): |
|
||||
l = self.limit |
|
||||
if l is None or l > 100: |
|
||||
r = 100 |
|
||||
else: |
|
||||
r = l |
|
||||
self.retrieve = r |
|
||||
return r > 0 |
|
||||
|
|
||||
async def fill_messages(self): |
|
||||
if not hasattr(self, 'channel'): |
|
||||
# do the required set up |
|
||||
channel = await self.messageable._get_channel() |
|
||||
self.channel = channel |
|
||||
|
|
||||
if self._get_retrieve(): |
|
||||
data = await self._retrieve_messages(self.retrieve) |
|
||||
if len(data) < 100: |
|
||||
self.limit = 0 # terminate the infinite loop |
|
||||
|
|
||||
if self.reverse: |
|
||||
data = reversed(data) |
|
||||
if self._filter: |
|
||||
data = filter(self._filter, data) |
|
||||
|
|
||||
channel = self.channel |
|
||||
for element in data: |
|
||||
await self.messages.put(self.state.create_message(channel=channel, data=element)) |
|
||||
|
|
||||
async def _retrieve_messages(self, retrieve) -> List[Message]: |
|
||||
"""Retrieve messages and update next parameters.""" |
|
||||
raise NotImplementedError |
|
||||
|
|
||||
async def _retrieve_messages_before_strategy(self, retrieve): |
|
||||
"""Retrieve messages using before parameter.""" |
|
||||
before = self.before.id if self.before else None |
|
||||
data: List[MessagePayload] = await self.logs_from(self.channel.id, retrieve, before=before) |
|
||||
if len(data): |
|
||||
if self.limit is not None: |
|
||||
self.limit -= retrieve |
|
||||
self.before = Object(id=int(data[-1]['id'])) |
|
||||
return data |
|
||||
|
|
||||
async def _retrieve_messages_after_strategy(self, retrieve): |
|
||||
"""Retrieve messages using after parameter.""" |
|
||||
after = self.after.id if self.after else None |
|
||||
data: List[MessagePayload] = await self.logs_from(self.channel.id, retrieve, after=after) |
|
||||
if len(data): |
|
||||
if self.limit is not None: |
|
||||
self.limit -= retrieve |
|
||||
self.after = Object(id=int(data[0]['id'])) |
|
||||
return data |
|
||||
|
|
||||
async def _retrieve_messages_around_strategy(self, retrieve): |
|
||||
"""Retrieve messages using around parameter.""" |
|
||||
if self.around: |
|
||||
around = self.around.id if self.around else None |
|
||||
data: List[MessagePayload] = await self.logs_from(self.channel.id, retrieve, around=around) |
|
||||
self.around = None |
|
||||
return data |
|
||||
return [] |
|
||||
|
|
||||
|
|
||||
class AuditLogIterator(_AsyncIterator['AuditLogEntry']): |
|
||||
def __init__(self, guild, limit=None, before=None, after=None, oldest_first=None, user_id=None, action_type=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 oldest_first is None: |
|
||||
self.reverse = after is not None |
|
||||
else: |
|
||||
self.reverse = oldest_first |
|
||||
|
|
||||
self.guild = guild |
|
||||
self.loop = guild._state.loop |
|
||||
self.request = guild._state.http.get_audit_logs |
|
||||
self.limit = limit |
|
||||
self.before = before |
|
||||
self.user_id = user_id |
|
||||
self.action_type = action_type |
|
||||
self.after = OLDEST_OBJECT |
|
||||
self._users = {} |
|
||||
self._state = guild._state |
|
||||
|
|
||||
self._filter = None # entry dict -> bool |
|
||||
|
|
||||
self.entries = asyncio.Queue() |
|
||||
|
|
||||
if self.reverse: |
|
||||
self._strategy = self._after_strategy |
|
||||
if self.before: |
|
||||
self._filter = lambda m: int(m['id']) < self.before.id |
|
||||
else: |
|
||||
self._strategy = self._before_strategy |
|
||||
if self.after and self.after != OLDEST_OBJECT: |
|
||||
self._filter = lambda m: int(m['id']) > self.after.id |
|
||||
|
|
||||
async def _before_strategy(self, retrieve): |
|
||||
before = self.before.id if self.before else None |
|
||||
data: AuditLogPayload = await self.request( |
|
||||
self.guild.id, limit=retrieve, user_id=self.user_id, action_type=self.action_type, before=before |
|
||||
) |
|
||||
|
|
||||
entries = data.get('audit_log_entries', []) |
|
||||
if len(data) and entries: |
|
||||
if self.limit is not None: |
|
||||
self.limit -= retrieve |
|
||||
self.before = Object(id=int(entries[-1]['id'])) |
|
||||
return data.get('users', []), entries |
|
||||
|
|
||||
async def _after_strategy(self, retrieve): |
|
||||
after = self.after.id if self.after else None |
|
||||
data: AuditLogPayload = await self.request( |
|
||||
self.guild.id, limit=retrieve, user_id=self.user_id, action_type=self.action_type, after=after |
|
||||
) |
|
||||
entries = data.get('audit_log_entries', []) |
|
||||
if len(data) and entries: |
|
||||
if self.limit is not None: |
|
||||
self.limit -= retrieve |
|
||||
self.after = Object(id=int(entries[0]['id'])) |
|
||||
return data.get('users', []), entries |
|
||||
|
|
||||
async def next(self) -> AuditLogEntry: |
|
||||
if self.entries.empty(): |
|
||||
await self._fill() |
|
||||
|
|
||||
try: |
|
||||
return self.entries.get_nowait() |
|
||||
except asyncio.QueueEmpty: |
|
||||
raise NoMoreItems() |
|
||||
|
|
||||
def _get_retrieve(self): |
|
||||
l = self.limit |
|
||||
if l is None or l > 100: |
|
||||
r = 100 |
|
||||
else: |
|
||||
r = l |
|
||||
self.retrieve = r |
|
||||
return r > 0 |
|
||||
|
|
||||
async def _fill(self): |
|
||||
from .user import User |
|
||||
|
|
||||
if self._get_retrieve(): |
|
||||
users, data = await self._strategy(self.retrieve) |
|
||||
if len(data) < 100: |
|
||||
self.limit = 0 # terminate the infinite loop |
|
||||
|
|
||||
if self.reverse: |
|
||||
data = reversed(data) |
|
||||
if self._filter: |
|
||||
data = filter(self._filter, data) |
|
||||
|
|
||||
for user in users: |
|
||||
u = User(data=user, state=self._state) |
|
||||
self._users[u.id] = u |
|
||||
|
|
||||
for element in data: |
|
||||
# TODO: remove this if statement later |
|
||||
if element['action_type'] is None: |
|
||||
continue |
|
||||
|
|
||||
await self.entries.put(AuditLogEntry(data=element, users=self._users, guild=self.guild)) |
|
||||
|
|
||||
|
|
||||
class GuildIterator(_AsyncIterator['Guild']): |
|
||||
"""Iterator for receiving the client's guilds. |
|
||||
|
|
||||
The guilds endpoint has the same two behaviours as described |
|
||||
in :class:`HistoryIterator`: |
|
||||
If ``before`` is specified, the guilds endpoint returns the ``limit`` |
|
||||
newest guilds before ``before``, sorted with newest first. For filling over |
|
||||
100 guilds, update the ``before`` parameter to the oldest guild received. |
|
||||
Guilds will be returned in order by time. |
|
||||
If `after` is specified, it returns the ``limit`` oldest guilds after ``after``, |
|
||||
sorted with newest first. For filling over 100 guilds, update the ``after`` |
|
||||
parameter to the newest guild received, If guilds are not reversed, they |
|
||||
will be out of order (99-0, 199-100, so on) |
|
||||
|
|
||||
Not that if both ``before`` and ``after`` are specified, ``before`` is ignored by the |
|
||||
guilds endpoint. |
|
||||
|
|
||||
Parameters |
|
||||
----------- |
|
||||
bot: :class:`discord.Client` |
|
||||
The client to retrieve the guilds from. |
|
||||
limit: :class:`int` |
|
||||
Maximum number of guilds to retrieve. |
|
||||
before: Optional[Union[:class:`abc.Snowflake`, :class:`datetime.datetime`]] |
|
||||
Object before which all guilds must be. |
|
||||
after: Optional[Union[:class:`abc.Snowflake`, :class:`datetime.datetime`]] |
|
||||
Object after which all guilds must be. |
|
||||
""" |
|
||||
|
|
||||
def __init__(self, bot, limit, before=None, after=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)) |
|
||||
|
|
||||
self.bot = bot |
|
||||
self.limit = limit |
|
||||
self.before = before |
|
||||
self.after = after |
|
||||
|
|
||||
self._filter = None |
|
||||
|
|
||||
self.state = self.bot._connection |
|
||||
self.get_guilds = self.bot.http.get_guilds |
|
||||
self.guilds = asyncio.Queue() |
|
||||
|
|
||||
if self.before and self.after: |
|
||||
self._retrieve_guilds = self._retrieve_guilds_before_strategy # type: ignore |
|
||||
self._filter = lambda m: int(m['id']) > self.after.id |
|
||||
elif self.after: |
|
||||
self._retrieve_guilds = self._retrieve_guilds_after_strategy # type: ignore |
|
||||
else: |
|
||||
self._retrieve_guilds = self._retrieve_guilds_before_strategy # type: ignore |
|
||||
|
|
||||
async def next(self) -> Guild: |
|
||||
if self.guilds.empty(): |
|
||||
await self.fill_guilds() |
|
||||
|
|
||||
try: |
|
||||
return self.guilds.get_nowait() |
|
||||
except asyncio.QueueEmpty: |
|
||||
raise NoMoreItems() |
|
||||
|
|
||||
def _get_retrieve(self): |
|
||||
l = self.limit |
|
||||
if l is None or l > 100: |
|
||||
r = 100 |
|
||||
else: |
|
||||
r = l |
|
||||
self.retrieve = r |
|
||||
return r > 0 |
|
||||
|
|
||||
def create_guild(self, data): |
|
||||
from .guild import Guild |
|
||||
|
|
||||
return Guild(state=self.state, data=data) |
|
||||
|
|
||||
async def fill_guilds(self): |
|
||||
if self._get_retrieve(): |
|
||||
data = await self._retrieve_guilds(self.retrieve) |
|
||||
if self.limit is None or len(data) < 100: |
|
||||
self.limit = 0 |
|
||||
|
|
||||
if self._filter: |
|
||||
data = filter(self._filter, data) |
|
||||
|
|
||||
for element in data: |
|
||||
await self.guilds.put(self.create_guild(element)) |
|
||||
|
|
||||
async def _retrieve_guilds(self, retrieve) -> List[Guild]: |
|
||||
"""Retrieve guilds and update next parameters.""" |
|
||||
raise NotImplementedError |
|
||||
|
|
||||
async def _retrieve_guilds_before_strategy(self, retrieve): |
|
||||
"""Retrieve guilds using before parameter.""" |
|
||||
before = self.before.id if self.before else None |
|
||||
data: List[GuildPayload] = await self.get_guilds(retrieve, before=before) |
|
||||
if len(data): |
|
||||
if self.limit is not None: |
|
||||
self.limit -= retrieve |
|
||||
self.before = Object(id=int(data[-1]['id'])) |
|
||||
return data |
|
||||
|
|
||||
async def _retrieve_guilds_after_strategy(self, retrieve): |
|
||||
"""Retrieve guilds using after parameter.""" |
|
||||
after = self.after.id if self.after else None |
|
||||
data: List[GuildPayload] = await self.get_guilds(retrieve, after=after) |
|
||||
if len(data): |
|
||||
if self.limit is not None: |
|
||||
self.limit -= retrieve |
|
||||
self.after = Object(id=int(data[0]['id'])) |
|
||||
return data |
|
||||
|
|
||||
|
|
||||
class MemberIterator(_AsyncIterator['Member']): |
|
||||
def __init__(self, guild, limit=1000, after=None): |
|
||||
|
|
||||
if isinstance(after, datetime.datetime): |
|
||||
after = Object(id=time_snowflake(after, high=True)) |
|
||||
|
|
||||
self.guild = guild |
|
||||
self.limit = limit |
|
||||
self.after = after or OLDEST_OBJECT |
|
||||
|
|
||||
self.state = self.guild._state |
|
||||
self.get_members = self.state.http.get_members |
|
||||
self.members = asyncio.Queue() |
|
||||
|
|
||||
async def next(self) -> Member: |
|
||||
if self.members.empty(): |
|
||||
await self.fill_members() |
|
||||
|
|
||||
try: |
|
||||
return self.members.get_nowait() |
|
||||
except asyncio.QueueEmpty: |
|
||||
raise NoMoreItems() |
|
||||
|
|
||||
def _get_retrieve(self): |
|
||||
l = self.limit |
|
||||
if l is None or l > 1000: |
|
||||
r = 1000 |
|
||||
else: |
|
||||
r = l |
|
||||
self.retrieve = r |
|
||||
return r > 0 |
|
||||
|
|
||||
async def fill_members(self): |
|
||||
if self._get_retrieve(): |
|
||||
after = self.after.id if self.after else None |
|
||||
data = await self.get_members(self.guild.id, self.retrieve, after) |
|
||||
if not data: |
|
||||
# no data, terminate |
|
||||
return |
|
||||
|
|
||||
if len(data) < 1000: |
|
||||
self.limit = 0 # terminate loop |
|
||||
|
|
||||
self.after = Object(id=int(data[-1]['user']['id'])) |
|
||||
|
|
||||
for element in reversed(data): |
|
||||
await self.members.put(self.create_member(element)) |
|
||||
|
|
||||
def create_member(self, data): |
|
||||
from .member import Member |
|
||||
|
|
||||
return Member(data=data, guild=self.guild, state=self.state) |
|
||||
|
|
||||
|
|
||||
class ArchivedThreadIterator(_AsyncIterator['Thread']): |
|
||||
def __init__( |
|
||||
self, |
|
||||
channel_id: int, |
|
||||
guild: Guild, |
|
||||
limit: Optional[int], |
|
||||
joined: bool, |
|
||||
private: bool, |
|
||||
before: Optional[Union[Snowflake, datetime.datetime]] = None, |
|
||||
): |
|
||||
self.channel_id = channel_id |
|
||||
self.guild = guild |
|
||||
self.limit = limit |
|
||||
self.joined = joined |
|
||||
self.private = private |
|
||||
self.http = guild._state.http |
|
||||
|
|
||||
if joined and not private: |
|
||||
raise ValueError('Cannot iterate over joined public archived threads') |
|
||||
|
|
||||
self.before: Optional[str] |
|
||||
if before is None: |
|
||||
self.before = None |
|
||||
elif isinstance(before, datetime.datetime): |
|
||||
if joined: |
|
||||
self.before = str(time_snowflake(before, high=False)) |
|
||||
else: |
|
||||
self.before = before.isoformat() |
|
||||
else: |
|
||||
if joined: |
|
||||
self.before = str(before.id) |
|
||||
else: |
|
||||
self.before = snowflake_time(before.id).isoformat() |
|
||||
|
|
||||
self.update_before: Callable[[ThreadPayload], str] = self.get_archive_timestamp |
|
||||
|
|
||||
if joined: |
|
||||
self.endpoint = self.http.get_joined_private_archived_threads |
|
||||
self.update_before = self.get_thread_id |
|
||||
elif private: |
|
||||
self.endpoint = self.http.get_private_archived_threads |
|
||||
else: |
|
||||
self.endpoint = self.http.get_public_archived_threads |
|
||||
|
|
||||
self.queue: asyncio.Queue[Thread] = asyncio.Queue() |
|
||||
self.has_more: bool = True |
|
||||
|
|
||||
async def next(self) -> Thread: |
|
||||
if self.queue.empty(): |
|
||||
await self.fill_queue() |
|
||||
|
|
||||
try: |
|
||||
return self.queue.get_nowait() |
|
||||
except asyncio.QueueEmpty: |
|
||||
raise NoMoreItems() |
|
||||
|
|
||||
@staticmethod |
|
||||
def get_archive_timestamp(data: ThreadPayload) -> str: |
|
||||
return data['thread_metadata']['archive_timestamp'] |
|
||||
|
|
||||
@staticmethod |
|
||||
def get_thread_id(data: ThreadPayload) -> str: |
|
||||
return data['id'] # type: ignore |
|
||||
|
|
||||
async def fill_queue(self) -> None: |
|
||||
if not self.has_more: |
|
||||
raise NoMoreItems() |
|
||||
|
|
||||
limit = 50 if self.limit is None else max(self.limit, 50) |
|
||||
data = await self.endpoint(self.channel_id, before=self.before, limit=limit) |
|
||||
|
|
||||
# This stuff is obviously WIP because 'members' is always empty |
|
||||
threads: List[ThreadPayload] = data.get('threads', []) |
|
||||
for d in reversed(threads): |
|
||||
self.queue.put_nowait(self.create_thread(d)) |
|
||||
|
|
||||
self.has_more = data.get('has_more', False) |
|
||||
if self.limit is not None: |
|
||||
self.limit -= len(threads) |
|
||||
if self.limit <= 0: |
|
||||
self.has_more = False |
|
||||
|
|
||||
if self.has_more: |
|
||||
self.before = self.update_before(threads[-1]) |
|
||||
|
|
||||
def create_thread(self, data: ThreadPayload) -> Thread: |
|
||||
from .threads import Thread |
|
||||
return Thread(guild=self.guild, state=self.guild._state, data=data) |
|
||||
Loading…
Reference in new issue