Initial Template.
This commit is contained in:
485
src/utils/ui/leo.py
Normal file
485
src/utils/ui/leo.py
Normal file
@@ -0,0 +1,485 @@
|
||||
from typing import List, Optional, Any, Dict
|
||||
import asyncio
|
||||
import logging
|
||||
import time
|
||||
from contextvars import copy_context, Context
|
||||
|
||||
import discord
|
||||
from discord.ui import Modal, View, Item
|
||||
|
||||
from meta.logger import log_action_stack, logging_context
|
||||
from meta.errors import SafeCancellation
|
||||
|
||||
from . import logger
|
||||
from ..lib import MessageArgs, error_embed
|
||||
|
||||
__all__ = (
|
||||
'LeoUI',
|
||||
'MessageUI',
|
||||
'LeoModal',
|
||||
'error_handler_for'
|
||||
)
|
||||
|
||||
|
||||
class LeoUI(View):
|
||||
"""
|
||||
View subclass for small-scale user interfaces.
|
||||
|
||||
While a 'View' provides an interface for managing a collection of components,
|
||||
a `LeoUI` may also manage a message, and potentially slave Views or UIs.
|
||||
The `LeoUI` also exposes more advanced cleanup and timeout methods,
|
||||
and preserves the context.
|
||||
"""
|
||||
|
||||
def __init__(self, *args, ui_name=None, context=None, **kwargs) -> None:
|
||||
super().__init__(*args, **kwargs)
|
||||
|
||||
if context is None:
|
||||
self._context = copy_context()
|
||||
else:
|
||||
self._context = context
|
||||
|
||||
self._name = ui_name or self.__class__.__name__
|
||||
self._context.run(log_action_stack.set, [*self._context[log_action_stack], self._name])
|
||||
|
||||
# List of slaved views to stop when this view stops
|
||||
self._slaves: List[View] = []
|
||||
|
||||
# TODO: Replace this with a substitutable ViewLayout class
|
||||
self._layout: Optional[tuple[tuple[Item, ...], ...]] = None
|
||||
|
||||
@property
|
||||
def _stopped(self) -> asyncio.Future:
|
||||
"""
|
||||
Return an future indicating whether the View has finished interacting.
|
||||
|
||||
Currently exposes a hidden attribute of the underlying View.
|
||||
May be reimplemented in future.
|
||||
"""
|
||||
return self._View__stopped
|
||||
|
||||
def to_components(self) -> List[Dict[str, Any]]:
|
||||
"""
|
||||
Extending component generator to apply the set _layout, if it exists.
|
||||
"""
|
||||
if self._layout is not None:
|
||||
# Alternative rendering using layout
|
||||
components = []
|
||||
for i, row in enumerate(self._layout):
|
||||
# Skip empty rows
|
||||
if not row:
|
||||
continue
|
||||
|
||||
# Since we aren't relying on ViewWeights, manually check width here
|
||||
if sum(item.width for item in row) > 5:
|
||||
raise ValueError(f"Row {i} of custom {self.__class__.__name__} is too wide!")
|
||||
|
||||
# Create the component dict for this row
|
||||
components.append({
|
||||
'type': 1,
|
||||
'components': [item.to_component_dict() for item in row]
|
||||
})
|
||||
else:
|
||||
components = super().to_components()
|
||||
|
||||
return components
|
||||
|
||||
def set_layout(self, *rows: tuple[Item, ...]) -> None:
|
||||
"""
|
||||
Set the layout of the rendered View as a matrix of items,
|
||||
or more precisely, a list of action rows.
|
||||
|
||||
This acts independently of the existing sorting with `_ViewWeights`,
|
||||
and overrides the sorting if applied.
|
||||
"""
|
||||
self._layout = rows
|
||||
|
||||
async def cleanup(self):
|
||||
"""
|
||||
Coroutine to run when timeing out, stopping, or cancelling.
|
||||
Generally cleans up any open resources, and removes any leftover components.
|
||||
"""
|
||||
logging.debug(f"{self!r} running default cleanup.", extra={'action': 'cleanup'})
|
||||
return None
|
||||
|
||||
def stop(self):
|
||||
"""
|
||||
Extends View.stop() to also stop all the slave views.
|
||||
Note that stopping is idempotent, so it is okay if close() also calls stop().
|
||||
"""
|
||||
for slave in self._slaves:
|
||||
slave.stop()
|
||||
super().stop()
|
||||
|
||||
async def close(self, msg=None):
|
||||
self.stop()
|
||||
await self.cleanup()
|
||||
|
||||
async def pre_timeout(self):
|
||||
"""
|
||||
Task to execute before actually timing out.
|
||||
This may cancel the timeout by refreshing or rescheduling it.
|
||||
(E.g. to ask the user whether they want to keep going.)
|
||||
|
||||
Default implementation does nothing.
|
||||
"""
|
||||
return None
|
||||
|
||||
async def on_timeout(self):
|
||||
"""
|
||||
Task to execute after timeout is complete.
|
||||
Default implementation calls cleanup.
|
||||
"""
|
||||
await self.cleanup()
|
||||
|
||||
async def __dispatch_timeout(self):
|
||||
"""
|
||||
This essentially extends View._dispatch_timeout,
|
||||
to include a pre_timeout task
|
||||
which may optionally refresh and hence cancel the timeout.
|
||||
"""
|
||||
if self._View__stopped.done():
|
||||
# We are already stopped, nothing to do
|
||||
return
|
||||
|
||||
with logging_context(action='Timeout'):
|
||||
try:
|
||||
await self.pre_timeout()
|
||||
except asyncio.TimeoutError:
|
||||
pass
|
||||
except asyncio.CancelledError:
|
||||
pass
|
||||
except Exception:
|
||||
logger.exception(
|
||||
"Unhandled error caught while dispatching timeout for {self!r}.",
|
||||
extra={'with_ctx': True, 'action': 'Error'}
|
||||
)
|
||||
|
||||
# Check if we still need to timeout
|
||||
if self.timeout is None:
|
||||
# The timeout was removed entirely, silently walk away
|
||||
return
|
||||
|
||||
if self._View__stopped.done():
|
||||
# We stopped while waiting for the pre timeout.
|
||||
# Or maybe another thread timed us out
|
||||
# Either way, we are done here
|
||||
return
|
||||
|
||||
now = time.monotonic()
|
||||
if self._View__timeout_expiry is not None and now < self._View__timeout_expiry:
|
||||
# The timeout was extended, make sure the timeout task is running then fade away
|
||||
if self._View__timeout_task is None or self._View__timeout_task.done():
|
||||
self._View__timeout_task = asyncio.create_task(self._View__timeout_task_impl())
|
||||
else:
|
||||
# Actually timeout, and call the post-timeout task for cleanup.
|
||||
self._really_timeout()
|
||||
await self.on_timeout()
|
||||
|
||||
def _dispatch_timeout(self):
|
||||
"""
|
||||
Overriding timeout method completely, to support interactive flow during timeout,
|
||||
and optional refreshing of the timeout.
|
||||
"""
|
||||
return self._context.run(asyncio.create_task, self.__dispatch_timeout())
|
||||
|
||||
def _really_timeout(self):
|
||||
"""
|
||||
Actuallly times out the View.
|
||||
This copies View._dispatch_timeout, apart from the `on_timeout` dispatch,
|
||||
which is now handled by `__dispatch_timeout`.
|
||||
"""
|
||||
if self._View__stopped.done():
|
||||
return
|
||||
|
||||
if self._View__cancel_callback:
|
||||
self._View__cancel_callback(self)
|
||||
self._View__cancel_callback = None
|
||||
|
||||
self._View__stopped.set_result(True)
|
||||
|
||||
def _dispatch_item(self, *args, **kwargs):
|
||||
"""Extending event dispatch to run in the instantiation context."""
|
||||
return self._context.run(super()._dispatch_item, *args, **kwargs)
|
||||
|
||||
async def on_error(self, interaction: discord.Interaction, error: Exception, item: Item):
|
||||
"""
|
||||
Default LeoUI error handle.
|
||||
This may be tail extended by subclasses to preserve the exception stack.
|
||||
"""
|
||||
try:
|
||||
raise error
|
||||
except SafeCancellation as e:
|
||||
if e.msg and not interaction.is_expired():
|
||||
try:
|
||||
if interaction.response.is_done():
|
||||
await interaction.followup.send(
|
||||
embed=error_embed(e.msg),
|
||||
ephemeral=True
|
||||
)
|
||||
else:
|
||||
await interaction.response.send_message(
|
||||
embed=error_embed(e.msg),
|
||||
ephemeral=True
|
||||
)
|
||||
except discord.HTTPException:
|
||||
pass
|
||||
logger.debug(
|
||||
f"Caught a safe cancellation from LeoUI: {e.details}",
|
||||
extra={'action': 'Cancel'}
|
||||
)
|
||||
except Exception:
|
||||
logger.exception(
|
||||
f"Unhandled interaction exception occurred in item {item!r} of LeoUI {self!r} from interaction: "
|
||||
f"{interaction.data}",
|
||||
extra={'with_ctx': True, 'action': 'UIError'}
|
||||
)
|
||||
# Explicitly handle the bugsplat ourselves
|
||||
splat = interaction.client.tree.bugsplat(interaction, error)
|
||||
await interaction.client.tree.error_reply(interaction, splat)
|
||||
|
||||
|
||||
class MessageUI(LeoUI):
|
||||
"""
|
||||
Simple single-message LeoUI, intended as a framework for UIs
|
||||
attached to a single interaction response.
|
||||
|
||||
UIs may also be sent as regular messages by using `send(channel)` instead of `run(interaction)`.
|
||||
"""
|
||||
|
||||
def __init__(self, *args, callerid: Optional[int] = None, **kwargs):
|
||||
super().__init__(*args, **kwargs)
|
||||
|
||||
# ----- UI state -----
|
||||
# User ID of the original caller (e.g. command author).
|
||||
# Mainly used for interaction usage checks and logging
|
||||
self._callerid = callerid
|
||||
|
||||
# Original interaction, if this UI is sent as an interaction response
|
||||
self._original: discord.Interaction = None
|
||||
|
||||
# Message holding the UI, when the UI is sent attached to a followup
|
||||
self._message: discord.Message = None
|
||||
|
||||
# Refresh lock, to avoid cache collisions on refresh
|
||||
self._refresh_lock = asyncio.Lock()
|
||||
|
||||
@property
|
||||
def channel(self):
|
||||
if self._original is not None:
|
||||
return self._original.channel
|
||||
else:
|
||||
return self._message.channel
|
||||
|
||||
# ----- UI API -----
|
||||
async def run(self, interaction: discord.Interaction, **kwargs):
|
||||
"""
|
||||
Run the UI as a response or followup to the given interaction.
|
||||
|
||||
Should be extended if more complex run mechanics are needed
|
||||
(e.g. registering listeners or setting up caches).
|
||||
"""
|
||||
await self.draw(interaction, **kwargs)
|
||||
|
||||
async def refresh(self, *args, thinking: Optional[discord.Interaction] = None, **kwargs):
|
||||
"""
|
||||
Reload and redraw this UI.
|
||||
|
||||
Primarily a hook-method for use by parents and other controllers.
|
||||
Performs a full data and reload and refresh (maintaining UI state, e.g. page n).
|
||||
"""
|
||||
async with self._refresh_lock:
|
||||
# Reload data
|
||||
await self.reload()
|
||||
# Redraw UI message
|
||||
await self.redraw(thinking=thinking)
|
||||
|
||||
async def quit(self):
|
||||
"""
|
||||
Quit the UI.
|
||||
|
||||
This usually involves removing the original message,
|
||||
and stopping or closing the underlying View.
|
||||
"""
|
||||
for child in self._slaves:
|
||||
# TODO: Better to use duck typing or interface typing
|
||||
if isinstance(child, MessageUI) and not child.is_finished():
|
||||
asyncio.create_task(child.quit())
|
||||
try:
|
||||
if self._original is not None and not self._original.is_expired():
|
||||
await self._original.delete_original_response()
|
||||
self._original = None
|
||||
if self._message is not None:
|
||||
await self._message.delete()
|
||||
self._message = None
|
||||
except discord.HTTPException:
|
||||
pass
|
||||
|
||||
# Note close() also runs cleanup and stop
|
||||
await self.close()
|
||||
|
||||
# ----- UI Flow -----
|
||||
async def interaction_check(self, interaction: discord.Interaction):
|
||||
"""
|
||||
Check the given interaction is authorised to use this UI.
|
||||
|
||||
Default implementation simply checks that the interaction is
|
||||
from the original caller.
|
||||
Extend for more complex logic.
|
||||
"""
|
||||
return interaction.user.id == self._callerid
|
||||
|
||||
async def make_message(self) -> MessageArgs:
|
||||
"""
|
||||
Create the UI message body, depening on the current state.
|
||||
|
||||
Called upon each redraw.
|
||||
Should handle caching if message construction is for some reason intensive.
|
||||
|
||||
Must be implemented by concrete UI subclasses.
|
||||
"""
|
||||
raise NotImplementedError
|
||||
|
||||
async def refresh_layout(self):
|
||||
"""
|
||||
Asynchronously refresh the message components,
|
||||
and explicitly set the message component layout.
|
||||
|
||||
Called just before redrawing, before `make_message`.
|
||||
|
||||
Must be implemented by concrete UI subclasses.
|
||||
"""
|
||||
raise NotImplementedError
|
||||
|
||||
async def reload(self):
|
||||
"""
|
||||
Reload and recompute the underlying data for this UI.
|
||||
|
||||
Must be implemented by concrete UI subclasses.
|
||||
"""
|
||||
raise NotImplementedError
|
||||
|
||||
async def draw(self, interaction, force_followup=False, **kwargs):
|
||||
"""
|
||||
Send the UI as a response or followup to the given interaction.
|
||||
|
||||
If the interaction has been responded to, or `force_followup` is set,
|
||||
creates a followup message instead of a response to the interaction.
|
||||
"""
|
||||
# Initial data loading
|
||||
await self.reload()
|
||||
# Set the UI layout
|
||||
await self.refresh_layout()
|
||||
# Fetch message arguments
|
||||
args = await self.make_message()
|
||||
|
||||
as_followup = force_followup or interaction.response.is_done()
|
||||
if as_followup:
|
||||
self._message = await interaction.followup.send(**args.send_args, **kwargs, view=self)
|
||||
else:
|
||||
self._original = interaction
|
||||
await interaction.response.send_message(**args.send_args, **kwargs, view=self)
|
||||
|
||||
async def send(self, channel: discord.abc.Messageable, **kwargs):
|
||||
"""
|
||||
Alternative to draw() which uses a discord.abc.Messageable.
|
||||
"""
|
||||
await self.reload()
|
||||
await self.refresh_layout()
|
||||
args = await self.make_message()
|
||||
self._message = await channel.send(**args.send_args, view=self)
|
||||
|
||||
async def _redraw(self, args):
|
||||
if self._original and not self._original.is_expired():
|
||||
await self._original.edit_original_response(**args.edit_args, view=self)
|
||||
elif self._message:
|
||||
await self._message.edit(**args.edit_args, view=self)
|
||||
else:
|
||||
# Interaction expired or already closed. Quietly cleanup.
|
||||
await self.close()
|
||||
|
||||
async def redraw(self, thinking: Optional[discord.Interaction] = None):
|
||||
"""
|
||||
Update the output message for this UI.
|
||||
|
||||
If a thinking interaction is provided, deletes the response while redrawing.
|
||||
"""
|
||||
await self.refresh_layout()
|
||||
args = await self.make_message()
|
||||
|
||||
if thinking is not None and not thinking.is_expired() and thinking.response.is_done():
|
||||
asyncio.create_task(thinking.delete_original_response())
|
||||
|
||||
try:
|
||||
await self._redraw(args)
|
||||
except discord.HTTPException as e:
|
||||
# Unknown communication error, nothing we can reliably do. Exit quietly.
|
||||
logger.warning(
|
||||
f"Unexpected UI redraw failure occurred in {self}: {repr(e)}",
|
||||
)
|
||||
await self.close()
|
||||
|
||||
async def cleanup(self):
|
||||
"""
|
||||
Remove message components from interaction response, if possible.
|
||||
|
||||
Extend to remove listeners or clean up caches.
|
||||
`cleanup` is always called when the UI is exiting,
|
||||
through timeout or user-driven closure.
|
||||
"""
|
||||
try:
|
||||
if self._original is not None and not self._original.is_expired():
|
||||
await self._original.edit_original_response(view=None)
|
||||
self._original = None
|
||||
if self._message is not None:
|
||||
await self._message.edit(view=None)
|
||||
self._message = None
|
||||
except discord.HTTPException:
|
||||
pass
|
||||
|
||||
|
||||
class LeoModal(Modal):
|
||||
"""
|
||||
Context-aware Modal class.
|
||||
"""
|
||||
def __init__(self, *args, context: Optional[Context] = None, **kwargs):
|
||||
super().__init__(**kwargs)
|
||||
|
||||
if context is None:
|
||||
self._context = copy_context()
|
||||
else:
|
||||
self._context = context
|
||||
self._context.run(log_action_stack.set, [*self._context[log_action_stack], self.__class__.__name__])
|
||||
|
||||
def _dispatch_submit(self, *args, **kwargs):
|
||||
"""
|
||||
Extending event dispatch to run in the instantiation context.
|
||||
"""
|
||||
return self._context.run(super()._dispatch_submit, *args, **kwargs)
|
||||
|
||||
def _dispatch_item(self, *args, **kwargs):
|
||||
"""Extending event dispatch to run in the instantiation context."""
|
||||
return self._context.run(super()._dispatch_item, *args, **kwargs)
|
||||
|
||||
async def on_error(self, interaction: discord.Interaction, error: Exception, *args):
|
||||
"""
|
||||
Default LeoModal error handle.
|
||||
This may be tail extended by subclasses to preserve the exception stack.
|
||||
"""
|
||||
try:
|
||||
raise error
|
||||
except Exception:
|
||||
logger.exception(
|
||||
f"Unhandled interaction exception occurred in {self!r}. Interaction: {interaction.data}",
|
||||
extra={'with_ctx': True, 'action': 'ModalError'}
|
||||
)
|
||||
# Explicitly handle the bugsplat ourselves
|
||||
splat = interaction.client.tree.bugsplat(interaction, error)
|
||||
await interaction.client.tree.error_reply(interaction, splat)
|
||||
|
||||
|
||||
def error_handler_for(exc):
|
||||
def wrapper(coro):
|
||||
coro._ui_error_handler_for_ = exc
|
||||
return coro
|
||||
return wrapper
|
||||
Reference in New Issue
Block a user