dolfies
3 years ago
11 changed files with 371 additions and 518 deletions
@ -1,306 +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 |
|
||||
from typing import Awaitable, TYPE_CHECKING, TypeVar, Optional, Any, Callable, Union, List, Tuple, AsyncIterator, Dict |
|
||||
|
|
||||
from .appinfo import InteractionApplication |
|
||||
from .errors import InvalidData |
|
||||
from .utils import _generate_nonce, _get_as_snowflake |
|
||||
from .object import Object |
|
||||
from .commands import _command_factory |
|
||||
from .enums import AppCommandType |
|
||||
|
|
||||
__all__ = ( |
|
||||
'CommandIterator', |
|
||||
'FakeCommandIterator', |
|
||||
) |
|
||||
|
|
||||
if TYPE_CHECKING: |
|
||||
from .user import User |
|
||||
from .message import Message |
|
||||
from .abc import Snowflake, Messageable |
|
||||
from .commands import ApplicationCommand |
|
||||
from .channel import DMChannel |
|
||||
|
|
||||
T = TypeVar('T') |
|
||||
OT = TypeVar('OT') |
|
||||
_Func = Callable[[T], Union[OT, Awaitable[OT]]] |
|
||||
|
|
||||
OLDEST_OBJECT = Object(id=0) |
|
||||
|
|
||||
|
|
||||
def _is_fake(item: Union[Messageable, Message]) -> bool: # I hate this too, but <circular imports> and performance exist |
|
||||
try: |
|
||||
item.guild # type: ignore |
|
||||
except AttributeError: |
|
||||
return True |
|
||||
try: |
|
||||
item.channel.me # type: ignore |
|
||||
except AttributeError: |
|
||||
return False |
|
||||
return True |
|
||||
|
|
||||
|
|
||||
class CommandIterator: |
|
||||
def __new__(cls, *args, **kwargs) -> Union[CommandIterator, FakeCommandIterator]: |
|
||||
if _is_fake(args[0]): |
|
||||
return FakeCommandIterator(*args) |
|
||||
else: |
|
||||
return super().__new__(cls) |
|
||||
|
|
||||
def __init__( |
|
||||
self, |
|
||||
item: Union[Messageable, Message], |
|
||||
type: AppCommandType, |
|
||||
query: Optional[str] = None, |
|
||||
limit: Optional[int] = None, |
|
||||
command_ids: Optional[List[int]] = None, |
|
||||
**kwargs, |
|
||||
) -> None: |
|
||||
if query and command_ids: |
|
||||
raise TypeError('Cannot specify both query and command_ids') |
|
||||
if limit is not None and limit <= 0: |
|
||||
raise ValueError('limit must be > 0') |
|
||||
|
|
||||
self.item = item |
|
||||
self.channel = None |
|
||||
self.state = item._state |
|
||||
self._tuple = None |
|
||||
self.type = type |
|
||||
_, self.cls = _command_factory(int(type)) |
|
||||
self.query = query |
|
||||
self.limit = limit |
|
||||
self.command_ids = command_ids |
|
||||
self.applications: bool = kwargs.get('applications', True) |
|
||||
self.application: Snowflake = kwargs.get('application', None) |
|
||||
self.commands = asyncio.Queue() |
|
||||
self._application_cache: Dict[int, InteractionApplication] = {} |
|
||||
|
|
||||
async def _process_args(self) -> Tuple[DMChannel, Optional[str], Optional[Union[User, Message]]]: |
|
||||
item = self.item |
|
||||
if self.type is AppCommandType.user: |
|
||||
channel = await item._get_channel() # type: ignore |
|
||||
if getattr(item, 'bot', None): |
|
||||
item = item |
|
||||
else: |
|
||||
item = None |
|
||||
text = 'user' |
|
||||
elif self.type is AppCommandType.message: |
|
||||
message = self.item |
|
||||
channel = message.channel # type: ignore |
|
||||
text = 'message' |
|
||||
elif self.type is AppCommandType.chat_input: |
|
||||
channel = await item._get_channel() # type: ignore |
|
||||
item = None |
|
||||
text = None |
|
||||
self._process_kwargs(channel) # type: ignore |
|
||||
return channel, text, item # type: ignore |
|
||||
|
|
||||
def _process_kwargs(self, channel) -> None: |
|
||||
kwargs = { |
|
||||
'guild_id': channel.guild.id, |
|
||||
'type': self.type.value, |
|
||||
'offset': 0, |
|
||||
} |
|
||||
if self.applications: |
|
||||
kwargs['applications'] = True # Only sent if it's True... |
|
||||
if app := self.application: |
|
||||
kwargs['application'] = app.id |
|
||||
if (query := self.query) is not None: |
|
||||
kwargs['query'] = query |
|
||||
if cmds := self.command_ids: |
|
||||
kwargs['command_ids'] = cmds |
|
||||
self.kwargs = kwargs |
|
||||
|
|
||||
async def iterate(self) -> AsyncIterator[ApplicationCommand]: |
|
||||
while True: |
|
||||
if self.commands.empty(): |
|
||||
await self.fill_commands() |
|
||||
|
|
||||
try: |
|
||||
yield self.commands.get_nowait() |
|
||||
except asyncio.QueueEmpty: |
|
||||
break |
|
||||
|
|
||||
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_commands(self) -> None: |
|
||||
if not self._tuple: # Do the required setup |
|
||||
self._tuple = await self._process_args() |
|
||||
|
|
||||
if not self._get_retrieve(): |
|
||||
return |
|
||||
|
|
||||
state = self.state |
|
||||
kwargs = self.kwargs |
|
||||
retrieve = self.retrieve |
|
||||
nonce = _generate_nonce() |
|
||||
|
|
||||
def predicate(d): |
|
||||
return d.get('nonce') == nonce |
|
||||
|
|
||||
data = None |
|
||||
for _ in range(3): |
|
||||
await state.ws.request_commands(**kwargs, limit=retrieve, nonce=nonce) |
|
||||
try: |
|
||||
data: Optional[Dict[str, Any]] = await asyncio.wait_for( |
|
||||
state.ws.wait_for('guild_application_commands_update', predicate), timeout=3 |
|
||||
) |
|
||||
except asyncio.TimeoutError: |
|
||||
pass |
|
||||
|
|
||||
if data is None: |
|
||||
raise InvalidData('Didn\'t receive a response from Discord') |
|
||||
|
|
||||
cmds = data['application_commands'] |
|
||||
if len(cmds) < retrieve: |
|
||||
self.limit = 0 |
|
||||
elif self.limit is not None: |
|
||||
self.limit -= retrieve |
|
||||
|
|
||||
kwargs['offset'] += retrieve |
|
||||
|
|
||||
for app in data.get('applications', []): |
|
||||
self._application_cache[int(app['id'])] = InteractionApplication(state=state, data=app) |
|
||||
|
|
||||
for cmd in cmds: |
|
||||
self.commands.put_nowait(self.create_command(cmd)) |
|
||||
|
|
||||
def create_command(self, data) -> ApplicationCommand: |
|
||||
channel, item, value = self._tuple # type: ignore |
|
||||
if item is not None: |
|
||||
kwargs = {item: value} |
|
||||
else: |
|
||||
kwargs = {} |
|
||||
app_id = _get_as_snowflake(data, 'application_id') |
|
||||
return self.cls(state=channel._state, data=data, channel=channel, application=self._application_cache.get(app_id), **kwargs) # type: ignore |
|
||||
|
|
||||
|
|
||||
class FakeCommandIterator: |
|
||||
def __init__( |
|
||||
self, |
|
||||
item: Union[User, Message, DMChannel], |
|
||||
type: AppCommandType, |
|
||||
query: Optional[str] = None, |
|
||||
limit: Optional[int] = None, |
|
||||
command_ids: Optional[List[int]] = None, |
|
||||
) -> None: |
|
||||
if query and command_ids: |
|
||||
raise TypeError('Cannot specify both query and command_ids') |
|
||||
if limit is not None and limit <= 0: |
|
||||
raise ValueError('limit must be > 0') |
|
||||
|
|
||||
self.item = item |
|
||||
self.channel = None |
|
||||
self._tuple = None |
|
||||
self.type = type |
|
||||
_, self.cls = _command_factory(int(type)) |
|
||||
self.query = query |
|
||||
self.limit = limit |
|
||||
self.command_ids = command_ids |
|
||||
self.has_more = False |
|
||||
self.commands = asyncio.Queue() |
|
||||
|
|
||||
async def _process_args(self) -> Tuple[DMChannel, Optional[str], Optional[Union[User, Message]]]: |
|
||||
item = self.item |
|
||||
if self.type is AppCommandType.user: |
|
||||
channel = await item._get_channel() # type: ignore |
|
||||
if getattr(item, 'bot', None): |
|
||||
item = item |
|
||||
else: |
|
||||
item = None |
|
||||
text = 'user' |
|
||||
elif self.type is AppCommandType.message: |
|
||||
message = self.item |
|
||||
channel = message.channel # type: ignore |
|
||||
text = 'message' |
|
||||
elif self.type is AppCommandType.chat_input: |
|
||||
channel = await item._get_channel() # type: ignore |
|
||||
item = None |
|
||||
text = None |
|
||||
if not channel.recipient.bot: # type: ignore # Type checker cannot understand this |
|
||||
raise TypeError('User is not a bot') |
|
||||
return channel, text, item # type: ignore |
|
||||
|
|
||||
async def iterate(self) -> AsyncIterator[ApplicationCommand]: |
|
||||
while True: |
|
||||
if self.commands.empty(): |
|
||||
await self.fill_commands() |
|
||||
|
|
||||
try: |
|
||||
yield self.commands.get_nowait() |
|
||||
except asyncio.QueueEmpty: |
|
||||
break |
|
||||
|
|
||||
async def fill_commands(self) -> None: |
|
||||
if self.has_more: |
|
||||
return |
|
||||
|
|
||||
if not (stuff := self._tuple): |
|
||||
self._tuple = channel, _, _ = await self._process_args() |
|
||||
else: |
|
||||
channel = stuff[0] |
|
||||
|
|
||||
limit = self.limit or -1 |
|
||||
data = await channel._state.http.get_application_commands(channel.recipient.id) |
|
||||
ids = self.command_ids |
|
||||
query = self.query and self.query.lower() |
|
||||
type = self.type.value |
|
||||
|
|
||||
for cmd in data: |
|
||||
if cmd['type'] != type: |
|
||||
continue |
|
||||
|
|
||||
if ids: |
|
||||
if not int(cmd['id']) in ids: |
|
||||
continue |
|
||||
|
|
||||
if query: |
|
||||
if not query in cmd['name'].lower(): |
|
||||
continue |
|
||||
|
|
||||
self.commands.put_nowait(self.create_command(cmd)) |
|
||||
limit -= 1 |
|
||||
if limit == 0: |
|
||||
break |
|
||||
|
|
||||
self.has_more = True |
|
||||
|
|
||||
def create_command(self, data) -> ApplicationCommand: |
|
||||
channel, item, value = self._tuple # type: ignore |
|
||||
if item is not None: |
|
||||
kwargs = {item: value} |
|
||||
else: |
|
||||
kwargs = {} |
|
||||
return self.cls(state=channel._state, data=data, channel=channel, **kwargs) |
|
||||
Loading…
Reference in new issue