|
@@ -0,0 +1,134 @@
|
|
|
|
|
+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,
|
|
|
|
|
+ *files_url: str,
|
|
|
|
|
+ save_path: Optional[str] = None) -> List[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
|
|
|
|
|
+ """
|
|
|
|
|
+
|
|
|
|
|
+ jobs = [self._job_factory(url, save_path=save_path)
|
|
|
|
|
+ for url in files_url]
|
|
|
|
|
+
|
|
|
|
|
+ tasks = [asyncio.ensure_future(job.download()) for job in jobs]
|
|
|
|
|
+
|
|
|
|
|
+ await asyncio.gather(*tasks)
|
|
|
|
|
+
|
|
|
|
|
+ return [task.result() for task in tasks]
|