From d6141878a3fb938fa0128560c572d58d443253b7 Mon Sep 17 00:00:00 2001 From: Conatum Date: Tue, 26 Apr 2022 13:45:01 +0300 Subject: [PATCH] (interactions): Add interaction manager. --- bot/meta/interactions/__init__.py | 1 + bot/meta/interactions/manager.py | 94 +++++++++++++++++++++++++++++++ 2 files changed, 95 insertions(+) create mode 100644 bot/meta/interactions/manager.py diff --git a/bot/meta/interactions/__init__.py b/bot/meta/interactions/__init__.py index 660b5a93..b522979a 100644 --- a/bot/meta/interactions/__init__.py +++ b/bot/meta/interactions/__init__.py @@ -2,3 +2,4 @@ from . import enums from .interactions import _component_interaction_factory, Interaction, ComponentInteraction, ModalResponse from .components import * from .modals import * +from .manager import InteractionManager diff --git a/bot/meta/interactions/manager.py b/bot/meta/interactions/manager.py new file mode 100644 index 00000000..7bdfab8b --- /dev/null +++ b/bot/meta/interactions/manager.py @@ -0,0 +1,94 @@ +import asyncio +from datetime import datetime, timedelta + + +class InteractionManager: + def __init__(self, timeout=600, extend=None): + self.futures = [] + self.self_futures = [] + + self.cleanup_function = self._cleanup + self.timeout_function = self._timeout + self.close_function = self._close + + self.timeout = timeout + self.extend = extend or timeout + self.expires_at = None + + self.cleaned_up = asyncio.Event() + + async def _timeout_loop(self): + diff = self.expires_at - datetime.now() + while True: + try: + await asyncio.sleep(diff) + except asyncio.CancelledError: + break + diff = self.expires_at - datetime.now() + if diff <= 0: + asyncio.create_task(self.timeout()) + break + + def extend_timeout(self): + new_expiry = max(datetime.now() + timedelta(seconds=self.extend), self.expires_at) + self.expires_at = new_expiry + + async def wait(self): + """ + Wait until the manager is "done". + That is, until all the futures are done, or `closed` is set. + """ + closed_task = asyncio.create_task(self.cleaned_up.wait()) + futures_task = asyncio.create_task(asyncio.wait(self.futures)) + await asyncio.wait((closed_task, futures_task), return_when=asyncio.FIRST_COMPLETED) + + async def __aenter__(self): + if self.timeout is not None: + self.expires_at = datetime.now() + timedelta(seconds=self.timeout) + self.self_futures.append(asyncio.create_task(self._timeout_loop())) + return self + + async def __aexit__(self, *args): + if not self.cleaned_up.is_set(): + await self.cleanup(exiting=True) + + async def _cleanup(self, manager, timeout=False, closing=False, exiting=False, **kwargs): + for future in self.futures: + future.cancel() + for future in self.self_futures: + future.cancel() + self.cleaned_up.set() + + def on_cleanup(self, func): + self.cleanup_function = func + return func + + async def cleanup(self, **kwargs): + await self.cleanup_function(self, **kwargs) + + async def _timeout(self, manager, **kwargs): + await self.cleanup(timeout=True, **kwargs) + + def on_timeout(self, func): + self.timeout_function = func + return func + + async def timeout(self): + await self.timeout_function(self) + + async def close(self, **kwargs): + """ + Request closure of the manager. + """ + await self.close_function(self, **kwargs) + + def on_close(self, func): + self.close_function = func + return func + + async def _close(self, manager, **kwargs): + await self.cleanup(closing=True, **kwargs) + + def add_future(self, future): + self.futures.append(future) + return future