itch-dl/itch_dl/api.py

from typing import Optional

import requests
from requests import Session
from urllib3.util.retry import Retry
from requests.adapters import HTTPAdapter

from .consts import ITCH_API


class ItchApiClient:
    def __init__(self, api_key: str, user_agent: str, base_url: Optional[str] = None):
        self.base_url = base_url or ITCH_API
        self.api_key = api_key

        self.requests = Session()
        self.requests.headers["User-Agent"] = user_agent

        retry_strategy = Retry(
            total=5,
            backoff_factor=10,
            allowed_methods=["HEAD", "GET"],
            status_forcelist=[429, 500, 502, 503, 504],
        )

        # No timeouts - set them explicitly on API calls below!
        adapter = HTTPAdapter(max_retries=retry_strategy)
        self.requests.mount("https://", adapter)
        self.requests.mount("http://", adapter)

    def get(
        self,
        endpoint: str,
        append_api_key: bool = True,
        guess_encoding: bool = False,
        **kwargs,
    ) -> requests.Response:
        """Wrapper around `requests.get`.

        :param endpoint: Path to fetch on the specified base URL.
        :param append_api_key: Send an authenticated API request.
        :param guess_encoding: Let requests guess the response encoding.
        """
        if append_api_key:
            params = kwargs.get("data") or {}

            if "api_key" not in params:
                params["api_key"] = self.api_key

            kwargs["data"] = params

        if endpoint.startswith("https://"):
            url = endpoint
        else:
            url = self.base_url + endpoint

        r = self.requests.get(url, **kwargs)

        # Itch always returns UTF-8 pages and API responses. Force
        # UTF-8 everywhere, except for binary file downloads.
        if not guess_encoding:
            r.encoding = "utf-8"

        return r