http_engine.py

import functools
import os
import asyncio
import time
import warnings
import logging
from typing import Optional, Iterable, Callable, AsyncIterator, Awaitable

import requests
import aiohttp
import nest_asyncio
from tqdm.auto import tqdm

DEFAULT_TIMEOUT = 20 * 60  # seconds
DEFAULT_CHUNKSIZE = 1024 * 10  # bytes


class HttpEngine:
    def __init__(
        self, 
        scheme: str, 
        host: str, 
        timeout: float = DEFAULT_TIMEOUT, 
        chunksize: int = DEFAULT_CHUNKSIZE,
        **kwargs
    ):
        """ Simple tool for making streaming http requests.

        Parameters
        ----------
        scheme :
            e.g "http" or "https"
        host : 
            will be used as the base for request urls
        timeout : 
            requests taking longer than this (in seconds) will raise a 
            `requests.Timeout` error. The clock on this timeout starts running 
            when the initial request is made.
        chunksize : 
            When streaming data, how many bytes ought to be requested at once.
        **kwargs : 
            unused. Defined here so that parameters can fall through from 
            subclasses
        """

        self.scheme = scheme
        self.host = host
        self.timeout = timeout
        self.chunksize = chunksize

    def _build_url(self, route):
        return f"{self.scheme}://{self.host}/{route}"

    def stream(self, route):
        """ Makes an http request and returns an iterator over the response.

        Parameters
        ----------
        route :
            the http route (under this object's host) to request against.

        """

        url = self._build_url(route)
        
        start_time = time.perf_counter()
        response = requests.get(url, stream=True)
        response_b = None
        if "Content-length" in response.headers:
            response_b = float(response.headers["Content-length"])

        size_message = f"{response_b / 1024 ** 2:3.3f}MiB" if response_b is not None else "potentially large"
        logging.warning(f"downloading a {size_message} file from {url}")
        progress = tqdm( unit="B", total=response_b, unit_scale=True,  desc="Downloading")

        for chunk in response.iter_content(self.chunksize):
            if chunk: # filter out keep-alive new chunks
                progress.update(len(chunk))
                yield chunk

            elapsed = time.perf_counter() - start_time
            if elapsed > self.timeout:
                raise requests.Timeout(f"Download took {elapsed} seconds, but timeout was set to {self.timeout}")

    @staticmethod
    def write_bytes(path: str, stream: Iterable[bytes]):
        write_from_stream(path, stream)


AsyncStreamCallbackType = Callable[[AsyncIterator[bytes]], Awaitable[None]]


class AsyncHttpEngine(HttpEngine):

    def __init__(
        self, 
        scheme: str, 
        host: str, 
        session: Optional[aiohttp.ClientSession] = None, 
        **kwargs
    ):
        """ Simple tool for making asynchronous streaming http requests.

        Parameters
        ----------
        scheme :
            e.g "http" or "https"
        host : 
            will be used as the base for request urls
        session : 
            If provided, this preconstructed session will be used rather than 
            a new one. Keep in mind that AsyncHttpEngine closes its session 
            when it is garbage collected!
        **kwargs :
            Will be passed to parent.

        """

        super(AsyncHttpEngine, self).__init__(scheme, host, **kwargs)

        self._session = None
        if session:
            self._session = session
            warnings.warn(
                "Recieved preconstructed session, ignoring timeout parameter."
            )

    @property
    def session(self):
        if self._session is None:
            self._session = aiohttp.ClientSession(
                timeout=aiohttp.client.ClientTimeout(self.timeout)
            )
        return self._session

    async def _stream_coroutine(
        self, 
        route: str, 
        callback: AsyncStreamCallbackType
    ):
        url = self._build_url(route)

        async with self.session.get(url) as response:
            await callback(response.content.iter_chunked(self.chunksize))

    def stream(
        self, 
        route: str
    ) -> Callable[[AsyncStreamCallbackType], Awaitable[None]]:
        """ Returns a coroutine which
            - makes an http request
            - exposes internally an asynchronous iterator over the response
            - takes a callback parameter, which should consume the iterator.

        Parameters
        ----------
        route :
            the http route (under this object's host) to request against.

        Notes
        -----
        To use this method, you will need an appropriate consumer. For
        instance, If you want to write the streamed data to a local file, you
        can use write_bytes_from_coroutine.

        Examples
        --------
        >>> engine = AsyncHttpEngine("http", "examplehost")
        >>> stream_coro = engine.stream("example/route")
        >>> write_bytes_from_coroutine("example/file/path.txt", stream_coro)

        """

        return functools.partial(self._stream_coroutine, route)

    def __del__(self):
        if getattr(self, "_session", None) is not None:
            nest_asyncio.apply()
            loop = asyncio.get_event_loop()
            loop.run_until_complete(self._session.close())

    @staticmethod
    def write_bytes(
            path: str,
            coroutine: Callable[[AsyncStreamCallbackType], Awaitable[None]]):
        write_bytes_from_coroutine(path, coroutine)


def write_bytes_from_coroutine(
    path: str, 
    coroutine: Callable[[AsyncStreamCallbackType], Awaitable[None]]
):
    """ Utility for streaming http from an asynchronous requester to a file.

    Parameters
    ----------
    path : 
        Write to this file
    coroutine : 
        The source of the data. Needs to have a specific structure, namely:
            - the first-position parameter of the coroutine ought to accept a
            callback. This callback ought to itself be awaitable.
            - within the coroutine, this callback ought to be called with a 
            single argument. That single argument should be an asynchronous 
            iterator.
        Please see AsyncHttpEngine.stream (and 
        AsyncHttpEngine._stream_coroutine) for an example. 
    
    """
    
    os.makedirs(os.path.dirname(path), exist_ok=True)
    
    async def callback(file_, iterable):
        async for chunk in iterable:
            file_.write(chunk)
            
    async def wrapper():
        with open(path, "wb") as file_:
            callback_ = functools.partial(callback, file_)
            await coroutine(callback_)

    nest_asyncio.apply()
    loop = asyncio.get_event_loop()
    loop.run_until_complete(wrapper())


def write_from_stream(path: str, stream: Iterable[bytes]):
    """ Write bytes to a file from an iterator

    Parameters
    ----------
    path : 
        write to this file
    stream : 
        iterable yielding bytes to be written

    """
    with open(path, "wb") as fil:
        for chunk in stream:
            fil.write(chunk)