From 4d1c0a77646e213e90bea85e7f30a1eda8c282f6 Mon Sep 17 00:00:00 2001
From: Maciej Olko <maciej.olko@affirm.com>
Date: Mon, 7 Aug 2023 08:21:39 +0200
Subject: [PATCH] Add throttling data to message and API (#456)

* Add throttling data to message and API

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* workaround a false positive typing check

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

---------

Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
---
 wlc/__init__.py  | 19 +++++++++++++++++--
 wlc/test_base.py |  6 +++++-
 wlc/test_wlc.py  |  5 ++++-
 3 files changed, 26 insertions(+), 4 deletions(-)

diff --git a/wlc/__init__.py b/wlc/__init__.py
index 13a2d6f..f2ee615 100644
--- a/wlc/__init__.py
+++ b/wlc/__init__.py
@@ -7,7 +7,7 @@
 import json
 import logging
 from copy import copy
-from typing import Any, ClassVar, Collection, Dict, Optional, Set, Tuple
+from typing import Any, ClassVar, Collection, Dict, Optional, Set, Tuple, cast
 from urllib.parse import urlencode, urlparse
 
 import dateutil.parser
@@ -37,6 +37,18 @@ def __init__(self, message: Optional[str] = None):
 class WeblateThrottlingError(WeblateException):
     """Throttling on the server."""
 
+    def __init__(self, limit: str, retry_after: str):
+        self.limit = limit
+        self.retry_after = retry_after
+        message_segments = [
+            cast(str, self.__doc__)
+        ]  # workaround for https://github.com/python/mypy/issues/15825
+        if limit:
+            message_segments.append(f"Limit is {limit} requests.")
+        if retry_after:
+            message_segments.append(f"Retry after {retry_after} seconds.")
+        super().__init__(" ".join(message_segments))
+
 
 class WeblatePermissionError(WeblateException):
     """You don't have permission to access this object."""
@@ -114,7 +126,10 @@ def process_error(self, error):
             status_code = error.response.status_code
 
             if status_code == 429:
-                raise WeblateThrottlingError
+                headers = error.response.headers
+                raise WeblateThrottlingError(
+                    headers.get("X-RateLimit-Limit"), headers.get("Retry-After")
+                )
             if status_code == 404:
                 raise WeblateException(
                     "Object not found on the server "
diff --git a/wlc/test_base.py b/wlc/test_base.py
index 43d5b08..7fb5646 100644
--- a/wlc/test_base.py
+++ b/wlc/test_base.py
@@ -207,7 +207,11 @@ def register_uris():
         method=responses.POST,
         json={"detail": "Can not create components"},
     )
-    register_error("projects/throttled", 429)
+    register_error(
+        "projects/throttled",
+        429,
+        headers={"X-RateLimit-Limit": "100", "Retry-After": "81818"},
+    )
     register_error("projects/error", 500)
     register_error("projects/io", 500, callback=raise_error)
     register_error("projects/bug", 500, callback=raise_error)
diff --git a/wlc/test_wlc.py b/wlc/test_wlc.py
index 504d763..39fa22c 100644
--- a/wlc/test_wlc.py
+++ b/wlc/test_wlc.py
@@ -50,7 +50,10 @@ def test_denied_json(self):
 
     def test_throttled(self):
         """Test listing projects."""
-        with self.assertRaisesRegex(WeblateException, "Throttling"):
+        with self.assertRaisesRegex(
+            WeblateException,
+            "Throttling.*Limit is 100 requests. Retry after 81818 seconds.",
+        ):
             Weblate().get_object("throttled")
 
     def test_error(self):