git.scottworley.com Git - paperdoorknob/blob - fetch.py

   1 # paperdoorknob: Print glowfic
   2 #
   3 # This program is free software: you can redistribute it and/or modify it
   4 # under the terms of the GNU General Public License as published by the
   5 # Free Software Foundation, version 3.
   6
   7
   8 from abc import ABC, abstractmethod
   9 from contextlib import contextmanager
  10 from typing import Iterator
  11
  12 import requests
  13 import requests_cache
  14
  15
  16 class Fetcher(ABC):
  17     @abstractmethod
  18     def fetch(self, url: str) -> bytes:
  19         raise NotImplementedError()
  20
  21
  22 class _SessionFetcher(Fetcher):
  23
  24     def __init__(self, session: requests.Session, timeout: int) -> None:
  25         self._session = session
  26         self._timeout = timeout
  27
  28     def fetch(self, url: str) -> bytes:
  29         with self._session.get(url, timeout=self._timeout) as r:
  30             r.raise_for_status()
  31             return r.content
  32
  33
  34 @contextmanager
  35 def DirectFetcher(timeout: int) -> Iterator[_SessionFetcher]:
  36     with requests.session() as session:
  37         yield _SessionFetcher(session, timeout)
  38
  39
  40 @contextmanager
  41 def CachingFetcher(cache_path: str, timeout: int) -> Iterator[_SessionFetcher]:
  42     with requests_cache.CachedSession(cache_path, cache_control=True) as session:
  43         yield _SessionFetcher(session, timeout)
  44
  45
  46 class FakeFetcher(Fetcher):
  47
  48     def __init__(self, resources: dict[str, bytes]) -> None:
  49         self._resources = resources
  50         self._fetch_count = 0
  51
  52     def fetch(self, url: str) -> bytes:
  53         self._fetch_count += 1
  54         if url not in self._resources:
  55             raise requests.HTTPError("URL not found")
  56         return self._resources[url]
  57
  58     def request_count(self) -> int:
  59         return self._fetch_count