Skip to content

feat: add reset_with_key and get_cooldown_time_with_key #1345

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
merged 8 commits into from
Apr 20, 2023
35 changes: 35 additions & 0 deletions interactions/models/internal/cooldowns.py
Original file line number Diff line number Diff line change
Expand Up @@ -387,6 +387,22 @@ async def get_cooldown_time(self, context: "BaseContext") -> float:
cooldown = await self.get_cooldown(context)
return cooldown.get_cooldown_time()

def get_cooldown_time_with_key(self, key: Any, *, create: bool = False) -> float:
"""
Get the remaining cooldown time with a key instead of the context.

Note:
The preferred way to get the cooldown system is to use `get_cooldown` as it will use the context to get the correct key.

Args:
key: The key to get the cooldown system for
create: Whether to create a new cooldown system if one does not exist
"""
cooldown = self.get_cooldown_with_key(key, create=create)
if cooldown is not None:
return cooldown.get_cooldown_time()
return 0

async def on_cooldown(self, context: "BaseContext") -> bool:
"""
Returns the cooldown state of the command.
Expand Down Expand Up @@ -423,6 +439,25 @@ async def reset(self, context: "BaseContext") -> None:
cooldown = await self.get_cooldown(context)
cooldown.reset()

def reset_with_key(self, key: Any) -> bool:
"""
Resets the cooldown for the bucket associated with the provided key.

Note:
The preferred way to reset the cooldown system is to use `reset_cooldown` as it will use the context to reset the correct cooldown.

Args:
key: The key to reset the cooldown system for

Returns:
True if the key existed and was reset successfully, False if the key didn't exist.
"""
cooldown = self.get_cooldown_with_key(key)
if cooldown is not None:
cooldown.reset()
return True
return False


class MaxConcurrency:
"""
Expand Down