TelegramBotAdmin/wbb/utils/aiodownloader/downloader.py

import os
import asyncio
from typing import Optional, List

import aiohttp
import aiofiles


class DownloadJob:
    """
    Download Job

    :param file_url: url where the file is located
    :param session: aiohttp session to be used on the job
    :param file_name: name to be used on the file. Defaults to the last part of the url
    :param save_path: dir where the file should be saved. Defaults to the current dir
    """

    def __init__(self,
                 session: aiohttp.ClientSession,
                 file_url: str,
                 save_path: Optional[str] = None,
                 chunk_size: Optional[int] = 1024):

        self.file_url = file_url
        self._session = session
        self._chunk_size = chunk_size

        self.file_name = file_url.split('/')[~0]
        self.file_path = os.path.join(
            save_path, self.file_name) if save_path else self.file_name

        self.completed = False
        self.progress = 0
        self.size = 0

    async def get_size(self) -> int:
        """
        Gets the content-length of the file from the file url.
        :return: the files size in bytes
        """
        if not self.size:
            async with self._session.get(self.file_url) as resp:
                if 200 <= resp.status < 300:
                    self.size = int(resp.headers['Content-Length'])

                else:
                    raise aiohttp.errors.HttpProcessingError(
                        message=f'There was a problem processing {self.file_url}',
                        code=resp.status)

        return self.size

    def _downloaded(self, chunk_size):
        """
        Method to be called when a chunk of the file is downloaded. It updates the
        progress, adds the size to the _advance variable and checks if the download
        is completed
        """
        self.progress += chunk_size

    async def download(self):
        """
        Downloads the file from the given url to a file in the given path.
        """

        async with self._session.get(self.file_url) as resp:
            # Checkning the response code
            if 200 <= resp.status < 300:
                # Saving the data to the file chunk by chunk.
                async with aiofiles.open(self.file_path, 'wb') as file:

                    # Downloading the file using the aiohttp.StreamReader
                    async for data in resp.content.iter_chunked(self._chunk_size):
                        await file.write(data)
                        self._downloaded(self._chunk_size)

                self.completed = True
                return self

            else:
                raise aiohttp.errors.HttpProcessingError(
                    message=f'There was a problem processing {self.file_url}', code=resp.status)


class Handler:
    """
    Top level interface with the downloader. It creates the download jobs and handles them.

    :param loop: asyncio loop. if not provided asyncio.get_event_loop() will be used
    :param session: aiohttp session to be used on all downloads. If not provided it will be
    created
    :param chunk_size: chunk bytes sizes to get from the file source. Defaults to 1024 bytes
    """

    def __init__(self,
                 loop: Optional[asyncio.BaseEventLoop] = None,
                 session: Optional[aiohttp.ClientSession] = None,
                 chunk_size: Optional[int] = 1024):

        self._loop = loop or asyncio.get_event_loop()
        self._session = session or aiohttp.ClientSession(loop=self._loop)
        self._chunk_size = chunk_size

    def _job_factory(self,
                     file_url: str,
                     save_path: Optional[str] = None) -> DownloadJob:
        """
        Shortcut for creating a download job. It adds the session and the chunk size.
        :param file_url: url where the file is located
        :param save_path: save path for the download
        :return:
        """
        return DownloadJob(self._session, file_url, save_path, self._chunk_size)

    async def download(self,
                       url: str,
                       save_path: Optional[str] = None) -> DownloadJob:
        """
        Downloads a bulk of files from the given list of urls to the given path.

        :param files_url: list of urls where the files are located
        :param save_path: path to be used for saving the files. Defaults to the current dir
        """

        job = self._job_factory(url, save_path=save_path)

        task = asyncio.ensure_future(job.download())

        await task
        file_name = url.split("/")[-1]
        path = os.getcwd() + "/" + file_name if not save_path else save_path
        return path