Add wait parameter to rate limits

docs/common.rst

@@ -3,6 +3,8 @@
 
 Common
 ======
+Common reference
+----------------
 
 .. autoclass:: Dispatcher
    :members:
@@ -17,3 +19,8 @@ Common
    :type: UndefinedType
 
    A alias for :attr:`UndefinedType.UNDEFINED`
+
+Errors
+------
+.. automodule:: nextcore.common.errors
+   :members:

nextcore/common/errors.py

@@ -0,0 +1,33 @@
+# The MIT License (MIT)
+# Copyright (c) 2021-present nextcore developers
+#
+# Permission is hereby granted, free of charge, to any person obtaining a
+# copy of this software and associated documentation files (the "Software"),
+# to deal in the Software without restriction, including without limitation
+# the rights to use, copy, modify, merge, publish, distribute, sublicense,
+# and/or sell copies of the Software, and to permit persons to whom the
+# Software is furnished to do so, subject to the following conditions:
+#
+# The above copyright notice and this permission notice shall be included in
+# all copies or substantial portions of the Software.
+#
+# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS
+# OR IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING
+# FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER
+# DEALINGS IN THE SOFTWARE.
+
+from __future__ import annotations
+
+from typing import TYPE_CHECKING
+
+if TYPE_CHECKING:
+    from typing import Final
+
+__all__: Final[tuple[str, ...]] = ()
+
+
+class RateLimitedError(Exception):
+    """A error for when a :class:`~nextcore.common.TimesPer` is rate limited and ``wait`` was :data:`False`"""

nextcore/common/times_per/times_per.py

@@ -27,6 +27,7 @@
 from queue import PriorityQueue
 from typing import TYPE_CHECKING, AsyncIterator
 
+from ..errors import RateLimitedError
 from .priority_queue_container import PriorityQueueContainer
 
 if TYPE_CHECKING:
@@ -59,14 +60,22 @@ def __init__(self, limit: int, per: float) -> None:
         self._pending_reset: bool = False
 
     @asynccontextmanager
-    async def acquire(self, *, priority: int = 0) -> AsyncIterator[None]:
+    async def acquire(self, *, priority: int = 0, wait: bool = True) -> AsyncIterator[None]:
         """Use a spot in the rate-limit.
 
         Parameters
         ----------
         priority:
             The priority. **Lower** number means it will be requested earlier.
+        wait:
+            Wait for a spot in the rate limit.
 
+            If this is :data:`False`, this will raise :exc:`RateLimitedError` instead.
+
+        Raises
+        ------
+        RateLimitedError
+            ``wait`` was set to :data:`False` and there was no more spots in the rate limit.
         Returns
         -------
         :class:`typing.AsyncContextManager`
@@ -78,6 +87,10 @@ async def acquire(self, *, priority: int = 0) -> AsyncIterator[None]:
         logger.debug("Reserved requests: %s", self._in_progress)
 
         if calculated_remaining == 0:
+            if not wait:
+                raise RateLimitedError()
+
+            # Wait for a spot
             future: Future[None] = Future()
             item = PriorityQueueContainer(priority, future)