91 lines
3.0 KiB
Python
91 lines
3.0 KiB
Python
import aiohttp
|
|
import asyncio
|
|
import aiofiles
|
|
|
|
import os
|
|
import re
|
|
|
|
from tqdm import tqdm
|
|
|
|
from urllib.parse import unquote
|
|
from mimetypes import guess_extension as extension
|
|
|
|
|
|
async def fetch_chunk(
|
|
session: aiohttp.ClientSession,
|
|
url: str,
|
|
start_byte: int,
|
|
end_byte: int,
|
|
part_num: int,
|
|
pbar: tqdm
|
|
):
|
|
headers = {'Range': f'bytes={start_byte}-{end_byte}'}
|
|
async with session.get(url, headers=headers) as response:
|
|
with open(f'part_{part_num}', 'wb') as f:
|
|
while True:
|
|
chunk = await response.content.read(1024)
|
|
if not chunk:
|
|
break
|
|
f.write(chunk)
|
|
pbar.update(len(chunk))
|
|
|
|
|
|
async def merge_files(filename, num_parts, pbar: tqdm):
|
|
pbar.set_description("Merging")
|
|
pbar.reset(total=pbar.total)
|
|
async with aiofiles.open(filename, 'wb') as f_out:
|
|
for i in range(num_parts):
|
|
async with aiofiles.open(f'part_{i}', 'rb') as f_in:
|
|
content = await f_in.read()
|
|
await f_out.write(content)
|
|
pbar.update(len(content))
|
|
os.remove(f'part_{i}')
|
|
|
|
|
|
def get_filename(response: aiohttp.ClientResponse):
|
|
headers = response.headers
|
|
|
|
if ("content-disposition" in headers
|
|
and "filename" in headers["content-disposition"]):
|
|
filename = re.match(
|
|
r'.*filename=\"(.+)\".*',
|
|
headers["content-disposition"]
|
|
).group(1)
|
|
return unquote(filename)
|
|
else:
|
|
url = str(response.url).split("?")[0]
|
|
filename = url.rstrip("/").split("/")[-1]
|
|
if re.findall(r'\.[a-zA-Z]{2}\w{0,2}$', filename):
|
|
return unquote(filename)
|
|
else:
|
|
content_type = headers["Content-Type"]
|
|
content_type = re.findall(r'([a-z]{4,11}/[\w+\-.]+)', content_type)[0]
|
|
if "Content-Type" in headers and extension(content_type):
|
|
filename += extension(content_type)
|
|
return unquote(filename)
|
|
else:
|
|
return unquote(filename)
|
|
|
|
|
|
async def download_file(url, num_parts=20):
|
|
async with aiohttp.ClientSession() as session:
|
|
async with session.head(url, allow_redirects=True) as response:
|
|
filename = get_filename(response)
|
|
content_length = int(response.headers['Content-Length'])
|
|
chunk_size = content_length // num_parts
|
|
|
|
with tqdm(total=content_length, unit="B", unit_scale=True,
|
|
desc="Downloading") as pbar:
|
|
tasks = []
|
|
for i in range(num_parts):
|
|
start_byte = i * chunk_size
|
|
if i == num_parts - 1:
|
|
end_byte = content_length
|
|
else:
|
|
end_byte = start_byte + chunk_size - 1
|
|
tasks.append(
|
|
fetch_chunk(session, url, start_byte, end_byte, i, pbar))
|
|
|
|
await asyncio.gather(*tasks)
|
|
await merge_files(filename, num_parts, pbar)
|