Add basic functions for scheduled state removal

This commit is contained in:
Luna 2020-04-04 16:12:03 -03:00
parent b67de9a62d
commit 720417d14d
1 changed files with 17 additions and 0 deletions

View File

@ -96,6 +96,8 @@ class StateManager:
#: raw mapping from session ids to GatewayState
self.states_raw = StateDictWrapper(self, {})
self.tasks = {}
def insert(self, state: GatewayState):
"""Insert a new state object."""
user_states = self.states[state.user_id]
@ -239,3 +241,18 @@ class StateManager:
# DMs and GDMs use all user states
return self.user_states(user_id)
async def _future_cleanup(self, state: GatewayState):
await asyncio.sleep(30)
self.remove(state)
async def schedule_deletion(self, state: GatewayState):
task = app.loop.create_task(self._future_cleanup(state))
self.tasks[state.session_id] = task
async def unschedule_deletion(self, state: GatewayState):
try:
task = self.tasks.pop(state.session_id)
task.cancel()
except KeyError:
pass