Cleaner Fetcher interface

[paperdoorknob] / fetch.py

# paperdoorknob: Print glowfic
#
# This program is free software: you can redistribute it and/or modify it
# under the terms of the GNU General Public License as published by the
# Free Software Foundation, version 3.


from abc import ABC, abstractmethod
from contextlib import contextmanager
from typing import Iterator

import requests
import requests_cache


class Fetcher(ABC):
    @abstractmethod
    def fetch(self, url: str) -> bytes:
        raise NotImplementedError()


class _SessionFetcher(Fetcher):

    def __init__(self, session: requests.Session, timeout: int) -> None:
        self._session = session
        self._timeout = timeout

    def fetch(self, url: str) -> bytes:
        with self._session.get(url, timeout=self._timeout) as r:
            r.raise_for_status()
            return r.content


@contextmanager
def DirectFetcher(timeout: int) -> Iterator[_SessionFetcher]:
    with requests.session() as session:
        yield _SessionFetcher(session, timeout)


@contextmanager
def CachingFetcher(cache_path: str, timeout: int) -> Iterator[_SessionFetcher]:
    with requests_cache.CachedSession(cache_path, cache_control=True) as session:
        yield _SessionFetcher(session, timeout)