mirror of
https://github.com/Nick80835/microbot
synced 2025-08-22 18:19:16 +00:00
79 lines
2.9 KiB
Python
79 lines
2.9 KiB
Python
import mimetypes
|
|
from asyncio import gather
|
|
from os import remove
|
|
|
|
import aiofiles
|
|
from aiohttp import ClientSession
|
|
|
|
|
|
class ParallelDownload:
|
|
def __init__(self, url: str, aioclient: ClientSession, file_name: str):
|
|
self.url = url
|
|
self.aioclient = aioclient
|
|
self.file_name = file_name
|
|
|
|
async def download_chunk(self, chunk_start: int, chunk_end: int, chunk_number: int) -> str:
|
|
async with self.aioclient.get(self.url, headers={"Range": f"bytes={chunk_start}-{chunk_end}"}) as response:
|
|
async with aiofiles.open(f"ubot/cache/{self.file_name}.part{chunk_number}", mode="wb") as cache_file:
|
|
while chunk := await response.content.read(4096):
|
|
await cache_file.write(chunk)
|
|
|
|
await cache_file.flush()
|
|
|
|
return f"ubot/cache/{self.file_name}.part{chunk_number}"
|
|
|
|
async def generate_chunk_coros(self, content_length: int, chunk_size: int) -> list:
|
|
place = 0
|
|
chunk_number = 0
|
|
chunk_coros = []
|
|
|
|
while place < content_length:
|
|
if place + chunk_size >= content_length:
|
|
chunk_coros.append(self.download_chunk(place, content_length, chunk_number))
|
|
break
|
|
|
|
chunk_coros.append(self.download_chunk(place, place + chunk_size, chunk_number))
|
|
|
|
place += chunk_size + 1
|
|
chunk_number += 1
|
|
|
|
return chunk_coros
|
|
|
|
async def get_download_info(self, threads) -> tuple[int, int, str|None]:
|
|
async with self.aioclient.get(self.url) as response:
|
|
content_length = int(response.headers["content-length"])
|
|
file_extension = mimetypes.guess_extension(response.headers["content-type"])
|
|
|
|
if threads:
|
|
if not content_length % threads:
|
|
chunk_size = content_length / threads
|
|
else:
|
|
chunk_size = content_length // (threads - 1)
|
|
else:
|
|
if content_length < 100000000:
|
|
chunk_size = 35000000
|
|
elif content_length < 500000000:
|
|
chunk_size = 100000000
|
|
else:
|
|
chunk_size = 300000000
|
|
|
|
return content_length, chunk_size, file_extension
|
|
|
|
|
|
async def download(url: str, file_name: str, aioclient: ClientSession = ClientSession(), threads: int|None = None) -> str:
|
|
downloader = ParallelDownload(url, aioclient, file_name)
|
|
content_length, chunk_size, file_extension = await downloader.get_download_info(threads)
|
|
chunk_coros = await downloader.generate_chunk_coros(content_length, chunk_size)
|
|
downloaded_part_files = await gather(*chunk_coros)
|
|
|
|
async with aiofiles.open(f"ubot/cache/{file_name}{file_extension}", "wb") as final_fh:
|
|
for part_file in downloaded_part_files:
|
|
async with aiofiles.open(part_file, "rb") as part_fh:
|
|
await final_fh.write(await part_fh.read())
|
|
|
|
remove(part_file)
|
|
|
|
await final_fh.flush()
|
|
|
|
return f"ubot/cache/{file_name}{file_extension}"
|