Skip to content

Commit

Permalink
rechore: BiliAPIClient.download
Browse files Browse the repository at this point in the history
  • Loading branch information
@Nemo2011
Nemo2011 committed Jan 30, 2025
1 parent e8fed56 commit 0e37617
Show file tree
Hide file tree
Showing 7 changed files with 280 additions and 83 deletions.
58 changes: 46 additions & 12 deletions bilibili_api/clients/AioHTTPClient.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -49,6 +49,8 @@ def __init__(
)
self.__wss: Dict[int, aiohttp.ClientWebSocketResponse] = {}
self.__ws_cnt: int = 0
self.__downloads: Dict[int, aiohttp.ClientResponse] = {}
self.__download_cnt: int = 0

def get_wrapped_session(self) -> aiohttp.ClientSession:
return self.__session
Expand Down Expand Up @@ -167,17 +169,47 @@ async def request(
)
return bili_api_resp

async def download(
self, url: str = "", headers: dict = {}, out: str = "", intro: str = "下载"
) -> None:
async with self.__session.get(url=url, headers=headers) as resp:
bts = 0
tot = resp.headers.get("content-length")
with open(out, "wb") as file:
async for chunk in resp.content.iter_chunks():
bts += file.write(chunk[0])
print(f"{intro} - {out} [{bts}/{tot}]", end="\r")
print()
async def download_create(
self,
url: str = "",
headers: dict = {},
) -> int:
if self.__need_update_session:
await self.__session.close()
self.__session = aiohttp.ClientSession(
loop=asyncio.get_event_loop(),
trust_env=self.__args["trust_env"],
connector=aiohttp.TCPConnector(verify_ssl=self.__args["verify_ssl"]),
)
self.__need_update_session = False
self.__download_cnt += 1
request_log.dispatch(
"DWN_CREATE",
"开始下载",
{
"id": self.__download_cnt,
"url": url,
"headers": headers,
},
)
self.__downloads[self.__download_cnt] = await self.__session.get(
url=url, headers=headers
)
return self.__download_cnt

async def download_chunk(self, cnt: int) -> bytes:
resp = self.__downloads[cnt]
data = await anext(resp.content.iter_chunked(1024))
request_log.dispatch(
"DWN_PART",
"收到部分下载数据",
{"id": cnt, "data": data},
)
return data

def download_content_length(self, cnt: int) -> int:
resp = self.__downloads[cnt]
return int(resp.headers.get("content-length", "0"))

async def ws_create(
self, url: str = "", params: dict = {}, headers: dict = {}
Expand Down Expand Up @@ -241,7 +273,9 @@ async def close(self):
set_verify_ssl.__doc__ = BiliAPIClient.set_verify_ssl.__doc__
set_trust_env.__doc__ = BiliAPIClient.set_trust_env.__doc__
request.__doc__ = BiliAPIClient.request.__doc__
download.__doc__ = BiliAPIClient.download.__doc__
download_create.__doc__ = BiliAPIClient.download_create.__doc__
download_chunk.__doc__ = BiliAPIClient.download_chunk.__doc__
download_content_length.__doc__ = BiliAPIClient.download_content_length.__doc__
ws_create.__doc__ = BiliAPIClient.ws_create.__doc__
ws_recv.__doc__ = BiliAPIClient.ws_recv.__doc__
ws_send.__doc__ = BiliAPIClient.ws_send.__doc__
Expand Down
52 changes: 39 additions & 13 deletions bilibili_api/clients/CurlCFFIClient.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -13,7 +13,7 @@
)
from curl_cffi import requests
import curl_cffi
from typing import Optional, Dict, Union, Tuple
from typing import Optional, Dict, Union, Tuple, AsyncGenerator
import asyncio


Expand Down Expand Up @@ -46,6 +46,8 @@ def __init__(
self.__ws_cnt: int = 0
self.__ws_need_close: Dict[int, bool] = {}
self.__ws_is_closed: Dict[int, bool] = {}
self.__downloads: Dict[int, requests.Response] = {}
self.__download_cnt: int = 0

def get_wrapped_session(self) -> requests.AsyncSession:
return self.__session
Expand Down Expand Up @@ -140,17 +142,39 @@ async def request(
)
return bili_api_resp

async def download(
self, url: str = "", headers: dict = {}, out: str = "", intro: str = "下载"
) -> None:
req = await self.__session.get(url=url, headers=headers, stream=True)
bts = 0
tot = req.headers.get("content-length")
with open(out, "wb") as file:
async for chunk in req.aiter_content():
bts += file.write(chunk)
print(f"{intro} - {out} [{bts}/{tot}]", end="\r")
print()
async def download_create(
self,
url: str = "",
headers: dict = {},
) -> int:
self.__download_cnt += 1
request_log.dispatch(
"DWN_CREATE",
"开始下载",
{
"id": self.__download_cnt,
"url": url,
"headers": headers,
},
)
self.__downloads[self.__download_cnt] = await self.__session.get(
url=url, headers=headers, stream=True
)
return self.__download_cnt

async def download_chunk(self, cnt: int) -> bytes:
resp = self.__downloads[cnt]
data = await anext(resp.aiter_content())
request_log.dispatch(
"DWN_PART",
"收到部分下载数据",
{"id": cnt, "data": data},
)
return data

def download_content_length(self, cnt: int) -> int:
resp = self.__downloads[cnt]
return int(resp.headers.get("content-length", "0"))

async def ws_create(
self, url: str = "", params: dict = {}, headers: dict = {}
Expand Down Expand Up @@ -241,7 +265,9 @@ async def close(self) -> None:
set_verify_ssl.__doc__ = BiliAPIClient.set_verify_ssl.__doc__
set_trust_env.__doc__ = BiliAPIClient.set_trust_env.__doc__
request.__doc__ = BiliAPIClient.request.__doc__
download.__doc__ = BiliAPIClient.download.__doc__
download_create.__doc__ = BiliAPIClient.download_create.__doc__
download_chunk.__doc__ = BiliAPIClient.download_chunk.__doc__
download_content_length.__doc__ = BiliAPIClient.download_content_length.__doc__
ws_create.__doc__ = BiliAPIClient.ws_create.__doc__
ws_recv.__doc__ = BiliAPIClient.ws_recv.__doc__
ws_send.__doc__ = BiliAPIClient.ws_send.__doc__
Expand Down
57 changes: 44 additions & 13 deletions bilibili_api/clients/HTTPXClient.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -12,7 +12,7 @@
)
from ..exceptions import ApiException
import httpx
from typing import Optional, Dict, Union
from typing import AsyncGenerator, Optional, Dict, Union


class HTTPXClient(BiliAPIClient):
Expand Down Expand Up @@ -41,6 +41,9 @@ def __init__(
verify=self.__verify_ssl,
trust_env=self.__trust_env,
)
self.__downloads: Dict[int, httpx.Response] = {}
self.__download_iter: Dict[int, AsyncGenerator] = {}
self.__download_cnt: int = 0

def get_wrapped_session(self) -> httpx.AsyncClient:
return self.__session
Expand Down Expand Up @@ -137,17 +140,43 @@ async def request(
)
return bili_api_resp

async def download(
self, url: str = "", headers: dict = {}, out: str = "", intro: str = "下载"
) -> None:
resp = await self.__session.get(url=url, headers=headers)
bts = 0
tot = resp.headers.get("content-length")
with open(out, "wb") as file:
async for chunk in resp.aiter_bytes():
bts += file.write(chunk)
print(f"{intro} - {out} [{bts}/{tot}]", end="\r")
print()
async def download_create(
self,
url: str = "",
headers: dict = {},
) -> int:
self.__download_cnt += 1
request_log.dispatch(
"DWN_CREATE",
"开始下载",
{
"id": self.__download_cnt,
"url": url,
"headers": headers,
},
)
req = self.__session.build_request(method="GET", url=url, headers=headers)
self.__downloads[self.__download_cnt] = await self.__session.send(
req, stream=True
)
self.__download_iter[self.__download_cnt] = self.__downloads[
self.__download_cnt
].aiter_bytes(1024)
return self.__download_cnt

async def download_chunk(self, cnt: int) -> bytes:
iter = self.__download_iter[cnt]
data = await anext(iter)
request_log.dispatch(
"DWN_PART",
"收到部分下载数据",
{"id": cnt, "data": data},
)
return data

def download_content_length(self, cnt: int) -> int:
resp = self.__downloads[cnt]
return int(resp.headers.get("content-length", "0"))

async def ws_create(self, *args, **kwargs) -> None:
"""
Expand Down Expand Up @@ -191,5 +220,7 @@ async def close(self) -> None:
set_verify_ssl.__doc__ = BiliAPIClient.set_verify_ssl.__doc__
set_trust_env.__doc__ = BiliAPIClient.set_trust_env.__doc__
request.__doc__ = BiliAPIClient.request.__doc__
download.__doc__ = BiliAPIClient.download.__doc__
download_create.__doc__ = BiliAPIClient.download_create.__doc__
download_chunk.__doc__ = BiliAPIClient.download_chunk.__doc__
download_content_length.__doc__ = BiliAPIClient.download_content_length.__doc__
close.__doc__ = BiliAPIClient.close.__doc__
Loading

0 comments on commit 0e37617

Please sign in to comment.