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