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 sys import stderr
  11 from typing import IO, Iterator
  12
  13 import requests
  14 import requests_cache
  15
  16
  17 _headers = {
  18     'User-Agent': 'paperdoorknob/0.0.1 (https://git.scottworley.com/paperdoorknob/)'}
  19
  20
  21 class Fetcher(ABC):
  22     @abstractmethod
  23     def fetch(self, url: str) -> bytes:
  24         raise NotImplementedError()
  25
  26
  27 class _SessionFetcher(Fetcher):
  28
  29     def __init__(self, session: requests.Session, timeout: int) -> None:
  30         self._session = session
  31         self._timeout = timeout
  32
  33     def fetch(self, url: str) -> bytes:
  34         with self._session.get(url, timeout=self._timeout, headers=_headers) as r:
  35             r.raise_for_status()
  36             return r.content
  37
  38
  39 class _CachingFetcher(Fetcher):
  40
  41     def __init__(
  42             self,
  43             session: requests_cache.CachedSession,
  44             timeout: int) -> None:
  45         self._session = session
  46         self._timeout = timeout
  47         self._request_count = 0
  48         self._cache_hit_count = 0
  49
  50     def fetch(self, url: str) -> bytes:
  51         with self._session.get(url, timeout=self._timeout, headers=_headers) as r:
  52             r.raise_for_status()
  53             self._request_count += 1
  54             self._cache_hit_count += int(r.from_cache)
  55             return r.content
  56
  57     @property
  58     def request_count(self) -> int:
  59         return self._request_count
  60
  61     @property
  62     def cache_hit_count(self) -> int:
  63         return self._cache_hit_count
  64
  65
  66 @contextmanager
  67 def DirectFetcher(timeout: int) -> Iterator[_SessionFetcher]:
  68     with requests.session() as session:
  69         yield _SessionFetcher(session, timeout)
  70
  71
  72 @contextmanager
  73 def CachingFetcher(
  74         cache_path: str,
  75         timeout: int,
  76         report_stream: IO[str] = stderr) -> Iterator[_CachingFetcher]:
  77     with requests_cache.CachedSession(cache_path, cache_control=True) as session:
  78         fetcher = _CachingFetcher(session, timeout)
  79         yield fetcher
  80         if fetcher.request_count > 0:
  81             percent = 100.0 * fetcher.cache_hit_count / fetcher.request_count
  82             print(
  83                 f"Fetch cache hits: {fetcher.cache_hit_count} ({percent:.1f}%)",
  84                 file=report_stream)
  85
  86
  87 class FakeFetcher(Fetcher):
  88
  89     def __init__(self, resources: dict[str, bytes]) -> None:
  90         self._resources = resources
  91         self._fetch_count = 0
  92
  93     def fetch(self, url: str) -> bytes:
  94         self._fetch_count += 1
  95         if url not in self._resources:
  96             raise requests.HTTPError("URL not found", url)
  97         return self._resources[url]
  98
  99     def request_count(self) -> int:
 100         return self._fetch_count