ratelimit.py

Go to the documentation of this file.

"""Ratelimit helper."""
 
from __future__ import annotations
 
import asyncio
from collections.abc import Callable, Hashable
import logging
import time
 
from homeassistant.core import HomeAssistant, callback
 
_LOGGER = logging.getLogger(__name__)
 
 
class KeyedRateLimit:
    """Class to track rate limits."""
 
    def __init__(
        self,
        hass: HomeAssistant,
    ) -> None:
        """Initialize ratelimit tracker."""
        self.hasshass = hass
        self._last_triggered: dict[Hashable, float] = {}
        self._rate_limit_timers: dict[Hashable, asyncio.TimerHandle] = {}
 
    @callback
    def async_has_timer(self, key: Hashable) -> bool:
        """Check if a rate limit timer is running."""
        return key in self._rate_limit_timers
 
    @callback
    def async_triggered(self, key: Hashable, now: float | None = None) -> None:
        """Call when the action we are tracking was triggered."""
        self.async_cancel_timerasync_cancel_timer(key)
        self._last_triggered[key] = now or time.time()
 
    @callback
    def async_cancel_timer(self, key: Hashable) -> None:
        """Cancel a rate limit time that will call the action."""
        if handle := self._rate_limit_timers.pop(key, None):
            handle.cancel()
 
    @callback
    def async_remove(self) -> None:
        """Remove all timers."""
        for timer in self._rate_limit_timers.values():
            timer.cancel()
        self._rate_limit_timers.clear()
 
    @callback
    def async_schedule_action[*_Ts](
        self,
        key: Hashable,
        rate_limit: float | None,
        now: float,
        action: Callable[[*_Ts], None],
        *args: *_Ts,
    ) -> float | None:
        """Check rate limits and schedule an action if we hit the limit.
 
        If the rate limit is hit:
            Schedules the action for when the rate limit expires
            if there are no pending timers. The action must
            be called in async.
 
            Returns the time the rate limit will expire
 
        If the rate limit is not hit:
 
            Return None
        """
        if rate_limit is None:
            return None
 
        if not (last_triggered := self._last_triggered.get(key)):
            return None
 
        next_call_time = last_triggered + rate_limit
 
        if next_call_time <= now:
            self.async_cancel_timerasync_cancel_timer(key)
            return None
 
        _LOGGER.debug(
            "Reached rate limit of %s for %s and deferred action until %s",
            rate_limit,
            key,
            next_call_time,
        )
 
        if key not in self._rate_limit_timers:
            self._rate_limit_timers[key] = self.hasshass.loop.call_later(
                next_call_time - now,
                action,
                *args,
            )
 
        return next_call_time