]> git.scottworley.com Git - paperdoorknob/blob - fetch.py
d5267534a100e1d0e527ef347547cb57b2ae6f31
[paperdoorknob] / 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