Merge pull request #9 from austind/develop

Implement support for HTTP-dates in `wait_from_header()`

Author: austind

README.md

@@ -23,8 +23,6 @@ Supports exceptions raised by both [`requests`](https://docs.python-requests.org
 
 ## Install
 
-## Install
-
 Install from PyPI:
 
 ```sh

docs/changelog.md

@@ -1,5 +1,10 @@
 # Changelog
 
+## v1.1.0
+
+* Add [HTTP-date](https://httpwg.org/specs/rfc9110.html#http.date) value parsing for [`retryhttp.wait_from_header`][]
+* [`is_rate_limited`][`retryhttp._utils.is_rate_limited`] now determines that a request was rate limited by the presence of a `Retry-After` header. Prior to v1.1.0, this was based on the status code `429 Too Many Requests`.
+
 ## v1.0.1
 
 * Fix documentation errors.

pyproject.toml

@@ -38,6 +38,7 @@ classifiers = [
 ]
 dependencies = [
     "httpx",
+    "pydantic",
     "requests",
     "tenacity"
 ]

retryhttp/_retry.py

@@ -46,7 +46,7 @@ def retry(
     """Retry potentially transient HTTP errors with sensible default behavior.
 
     By default, retries the following errors, for a total of 3 attempts, with
-    exponential backoff (except for `429 Too Many Requests`, which defaults to the
+    exponential backoff (except when rate limited, which defaults to the
     `Retry-After` header, if present):
 
     - HTTP status errors:
@@ -71,11 +71,11 @@ def retry(
         retry_server_errors: Whether to retry 5xx server errors.
         retry_network_errors: Whether to retry network errors.
         retry_timeouts: Whether to retry timeouts.
-        retry_rate_limited: Whether to retry `429 Too Many Requests` errors.
+        retry_rate_limited: Whether to retry when `Retry-After` header received.
         wait_server_errors: Wait strategy to use for server errors.
         wait_network_errors: Wait strategy to use for network errors.
         wait_timeouts: Wait strategy to use for timeouts.
-        wait_rate_limited: Wait strategy to use for `429 Too Many Requests` errors.
+        wait_rate_limited: Wait strategy to use when `Retry-After` header received.
         server_error_codes: One or more 5xx error codes that will trigger `wait_server_errors`
             if `retry_server_errors` is `True`. Defaults to 500, 502, 503, and 504.
         network_errors: One or more exceptions that will trigger `wait_network_errors` if
@@ -85,6 +85,7 @@ def retry(
             - `httpx.ReadError`
             - `httpx.WriteError`
             - `requests.ConnectError`
+            - `requests.exceptions.ChunkedEncodingError`
         timeouts: One or more exceptions that will trigger `wait_timeouts` if
             `retry_timeouts` is `True`. Defaults to:
 
@@ -95,7 +96,7 @@ def retry(
         Decorated function.
 
     Raises:
-        RuntimeError: if `retry_server_errors`, `retry_network_errors`, `retry_timeouts`,
+        RuntimeError: If `retry_server_errors`, `retry_network_errors`, `retry_timeouts`,
             and `retry_rate_limited` are all `False`.
 
     """
@@ -167,7 +168,7 @@ def __init__(
 
 
 class retry_if_rate_limited(retry_base):
-    """Retry if server responds with `429 Too Many Requests` (rate limited)."""
+    """Retry if server responds with a `Retry-After` header."""
 
     def __call__(self, retry_state: RetryCallState) -> bool:
         if retry_state.outcome and retry_state.outcome.failed:

retryhttp/_types.py

@@ -1,4 +1,20 @@
+from datetime import datetime
 from typing import Any, Callable, TypeVar
 
+
+class HTTPDate(str):
+    @classmethod
+    def __get_validators__(cls):
+        yield cls.validate
+
+    @classmethod
+    def validate(cls, value: str) -> str:
+        try:
+            datetime.strptime(value, "%a, %d %b %Y %H:%M:%S GMT")
+        except ValueError:
+            raise ValueError(f"Invalid HTTP-date format: {value}")
+        return value
+
+
 F = TypeVar("F", bound=Callable[..., Any])
 WrappedFn = TypeVar("WrappedFn", bound=Callable[..., Any])

retryhttp/_utils.py

@@ -1,7 +1,7 @@
+from datetime import datetime, timedelta, timezone
 from typing import Optional, Sequence, Tuple, Type, Union
 
-import httpx
-import requests
+from retryhttp._types import HTTPDate
 
 _HTTPX_INSTALLED = False
 _REQUESTS_INSTALLED = False
@@ -60,14 +60,8 @@ def get_default_timeouts() -> (
 ):
     """Get all timeout exceptions to use by default.
 
-    Args:
-        N/A
-
     Returns:
-        Tuple of timeout exceptions.
-
-    Raises:
-        N/A
+        tuple: Timeout exceptions.
 
     """
     exceptions = []
@@ -83,14 +77,8 @@ def get_default_http_status_exceptions() -> (
 ):
     """Get default HTTP status 4xx or 5xx exceptions.
 
-    Args:
-        N/A
-
     Returns:
-        Tuple of HTTP status exceptions.
-
-    Raises:
-        N/A
+        tuple: HTTP status exceptions.
 
     """
     exceptions = []
@@ -102,20 +90,22 @@ def get_default_http_status_exceptions() -> (
 
 
 def is_rate_limited(exc: Union[BaseException, None]) -> bool:
-    """Whether a given exception indicates a 429 Too Many Requests error.
+    """Whether a given exception indicates the user has been rate limited.
+
+    Rate limiting should return a `429 Too Many Requests` status, but in
+    practice, servers may return `503 Service Unavailable`, or possibly
+    another code. In any case, if rate limiting is the issue, the server
+    will include a `Retry-After` header.
 
     Args:
         exc: Exception to consider.
 
     Returns:
-        Boolean of whether exc indicates a 429 Too Many Requests error.
-
-    Raises:
-        N/A
+        bool: Whether exc indicates rate limiting.
 
     """
     if isinstance(exc, get_default_http_status_exceptions()):
-        return exc.response.status_code == 429
+        return "retry-after" in exc.response.headers.keys()
     return False
 
 
@@ -131,14 +121,31 @@ def is_server_error(
             to all (500-599).
 
     Returns:
-        Boolean of whether exc indicates an error included in status_codes.
-
-    Raises:
-        N/A
+        bool: whether exc indicates an error included in status_codes.
 
     """
     if isinstance(status_codes, int):
         status_codes = [status_codes]
     if isinstance(exc, get_default_http_status_exceptions()):
         return exc.response.status_code in status_codes
     return False
+
+
+def get_http_date(delta_seconds: int = 0) -> HTTPDate:
+    """Builds a valid HTTP-date string.
+
+    By default, returns an HTTP-date string for the current timestamp.
+
+    Args:
+        delta_seconds (int): Number of seconds to offset the timestamp
+            by. If a negative integer is passed, result will be in the
+            past.
+
+    Returns:
+        HTTPDate: A valid HTTP-date string.
+
+    """
+    date = datetime.now(timezone.utc)
+    if delta_seconds:
+        date = date + timedelta(seconds=delta_seconds)
+    return HTTPDate(date.strftime("%a, %d %b %Y %H:%M:%S GMT"))

retryhttp/_wait.py

@@ -1,3 +1,4 @@
+from datetime import datetime, timezone
 from typing import Sequence, Tuple, Type, Union
 
 from tenacity import RetryCallState, wait_exponential, wait_random_exponential
@@ -15,6 +16,13 @@
 class wait_from_header(wait_base):
     """Wait strategy that derives the wait value from an HTTP header.
 
+    Value may be either an integer representing the number of seconds to wait
+    before retrying, or a date in HTTP-date format, indicating when it is
+    acceptable to retry the request. If such a date value is found, this method
+    will use that value to determine the correct number of seconds to wait.
+
+    More info on HTTP-date format: https://httpwg.org/specs/rfc9110.html#http.date
+
     Args:
         header: Header to attempt to derive wait value from.
         fallback: Wait strategy to use if `header` is not present, or unable
@@ -34,21 +42,31 @@ def __call__(self, retry_state: RetryCallState) -> float:
         if retry_state.outcome:
             exc = retry_state.outcome.exception()
             if isinstance(exc, get_default_http_status_exceptions()):
+                value = exc.response.headers.get(self.header, "")
+
                 try:
-                    return float(
-                        exc.response.headers.get(
-                            self.header, self.fallback(retry_state)
-                        )
-                    )
+                    return float(value)
                 except ValueError:
                     pass
+
+                try:
+                    retry_after = datetime.strptime(value, "%a, %d %b %Y %H:%M:%S GMT")
+                    retry_after = retry_after.replace(tzinfo=timezone.utc)
+                    now = datetime.now(timezone.utc)
+                    return float((retry_after - now).seconds)
+                except ValueError:
+                    pass
+
         return self.fallback(retry_state)
 
 
 class wait_rate_limited(wait_from_header):
-    """Wait strategy to use when the server responds with `429 Too Many Requests`.
+    """Wait strategy to use when the server responds with a `Retry-After` header.
 
-    Attempts to derive wait value from the `Retry-After` header.
+    The `Retry-After` header may be sent with the `503 Service Unavailable` or
+    `429 Too Many Requests` status code. The header value may provide a date for when
+    you may retry the request, or an integer, indicating the number of seconds
+    to wait before retrying.
 
     Args:
         fallback: Wait strategy to use if `Retry-After` header is not present, or unable
@@ -70,7 +88,7 @@ class wait_context_aware(wait_base):
         wait_server_errors: Wait strategy to use with server errors.
         wait_network_errors: Wait strategy to use with network errors.
         wait_timeouts: Wait strategy to use with timeouts.
-        wait_rate_limited: Wait strategy to use with `429 Too Many Requests`.
+        wait_rate_limited: Wait strategy to use when rate limited.
         server_error_codes: One or more 5xx HTTP status codes that will trigger
             `wait_server_errors`.
         network_errors: One or more exceptions that will trigger `wait_network_errors`.
@@ -88,6 +106,7 @@ class wait_context_aware(wait_base):
             - `httpx.ReadTimeout`
             - `httpx.WriteTimeout`
             - `requests.Timeout`
+
     """
 
     def __init__(

tests/test_rate_limited.py

@@ -1,14 +1,19 @@
+from typing import Union
+
 import httpx
 import pytest
 import respx
+from pydantic import PositiveInt
 from tenacity import RetryError, retry, stop_after_attempt
 
 from retryhttp import retry_if_rate_limited, wait_rate_limited
+from retryhttp._types import HTTPDate
+from retryhttp._utils import get_http_date
 
 MOCK_URL = "https://example.com/"
 
 
-def rate_limited_response(retry_after: int = 1):
+def rate_limited_response(retry_after: Union[HTTPDate, PositiveInt] = 1):
     return httpx.Response(
         status_code=httpx.codes.TOO_MANY_REQUESTS,
         headers={"Retry-After": str(retry_after)},
@@ -30,9 +35,24 @@ def retry_rate_limited():
 def test_rate_limited_failure():
     route = respx.get(MOCK_URL).mock(
         side_effect=[
-            rate_limited_response(),
-            rate_limited_response(),
-            rate_limited_response(),
+            rate_limited_response(retry_after=1),
+            rate_limited_response(retry_after=1),
+            rate_limited_response(retry_after=1),
+        ]
+    )
+    with pytest.raises(RetryError):
+        retry_rate_limited()
+    assert route.call_count == 3
+    assert route.calls[2].response.status_code == 429
+
+
+@respx.mock
+def test_rate_limited_failure_httpdate():
+    route = respx.get(MOCK_URL).mock(
+        side_effect=[
+            rate_limited_response(retry_after=get_http_date(delta_seconds=2)),
+            rate_limited_response(retry_after=get_http_date(delta_seconds=2)),
+            rate_limited_response(retry_after=get_http_date(delta_seconds=2)),
         ]
     )
     with pytest.raises(RetryError):