ask_candid/base/api_base.py

from typing import Any

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


class BaseAPI:

    def __init__(
        self,
        url: str,
        headers: dict[str, Any] | None = None,
        total_retries: int = 3,
        backoff_factor: int = 2
    ) -> None:
        total_retries = max(total_retries, 10)

        adapter = HTTPAdapter(
            max_retries=Retry(
                total=total_retries,
                status_forcelist=[429, 500, 502, 503, 504],
                allowed_methods=frozenset({"HEAD", "GET", "POST", "OPTIONS"}),
                backoff_factor=backoff_factor,
            )
        )
        self.session = requests.Session()
        self.session.mount("https://", adapter)
        self.session.mount("http://", adapter)

        self.__url = url
        self.__headers = headers

    def get(self, **request_kwargs):
        r = self.session.get(url=self.__url, headers=self.__headers, params=request_kwargs, timeout=30)
        r.raise_for_status()
        return r.json()

    def post(self, payload: dict[str, Any]):
        r = self.session.post(url=self.__url, headers=self.__headers, json=payload, timeout=30)
        r.raise_for_status()
        return r.json()