diff --git a/biliarchiver/cli_tools/get_command.py b/biliarchiver/cli_tools/get_command.py index 3ab5383..25cfcff 100644 --- a/biliarchiver/cli_tools/get_command.py +++ b/biliarchiver/cli_tools/get_command.py @@ -15,7 +15,10 @@ from rich import print from biliarchiver.i18n import _, ngettext -async def by_series(url_or_sid: str) -> Path: +async def by_series(url_or_sid: str, truncate: int = int(1e10)) -> Path: + """ + truncate: do noting + """ sid = sid = ( re.search(r"sid=(\d+)", url_or_sid).groups()[0] if url_or_sid.startswith("http") @@ -78,7 +81,7 @@ def by_ranking(rid: int) -> Path: return Path(abs_filepath) -async def by_up_videos(url_or_mid: str) -> Path: +async def by_up_videos(url_or_mid: str, truncate: int = int(1e10)) -> Path: """频率高了会封""" if isinstance(url_or_mid, int): @@ -127,6 +130,10 @@ async def by_up_videos(url_or_mid: str) -> Path: _x, _y, bv_ids_page = await api.get_up_video_info(client, mid, pn, ps, order, keyword) bv_ids += bv_ids_page + if len(bv_ids) >= truncate: + print("truncate at", truncate) + break + print(mid, up_name, total_size) await client.aclose() assert len(bv_ids) == len(set(bv_ids)), _("有重复的 bv_id") @@ -212,7 +219,7 @@ def not_got_popular_series() -> list[int]: return series_not_got -async def by_favlist(url_or_fid: str): +async def by_favlist(url_or_fid: str, truncate: int = int(1e10)) -> Path: if url_or_fid.startswith("http"): fid = re.findall(r"fid=(\d+)", url_or_fid)[0] else: @@ -239,6 +246,11 @@ async def by_favlist(url_or_fid: str): ), end="\r", ) + + if len(bvids) >= truncate: + print("truncate at", truncate) + break + await asyncio.sleep(2) page_num += 1 await client.aclose() @@ -263,6 +275,8 @@ async def by_favlist(url_or_fid: str): ) ) + return Path(abs_filepath) + async def main( series: str, diff --git a/biliarchiver/cli_tools/utils.py b/biliarchiver/cli_tools/utils.py index 6ffde3b..4f6586c 100644 --- a/biliarchiver/cli_tools/utils.py +++ b/biliarchiver/cli_tools/utils.py @@ -1,4 +1,5 @@ from pathlib import Path +from typing import List, Union from biliarchiver.utils.identifier import is_bvid from biliarchiver.i18n import _ @@ -21,3 +22,8 @@ def read_bvids(bvids: str) -> list[str]: assert bvids_list is not None and len(bvids_list) > 0, _("bvids 为空") return bvids_list + +def read_bvids_from_txt(txt_path: Union[Path,str]) -> List[str]: + with open(txt_path, "r", encoding="utf-8") as f: + bvids = [line.strip() for line in f if line.strip().startswith("BV")] + return bvids \ No newline at end of file diff --git a/biliarchiver/rest_api/main.py b/biliarchiver/rest_api/main.py index b6ccf01..08cc4f0 100644 --- a/biliarchiver/rest_api/main.py +++ b/biliarchiver/rest_api/main.py @@ -1,17 +1,21 @@ import asyncio -from asyncio import Queue from contextlib import asynccontextmanager from datetime import datetime +import os +from pathlib import Path from typing import List +from biliarchiver.cli_tools.utils import read_bvids_from_txt + try: - from fastapi import FastAPI + from fastapi import FastAPI, HTTPException except ImportError: print("Please install fastapi") print("`pip install fastapi`") raise +from biliarchiver.cli_tools.get_command import by_favlist, by_series, by_up_videos from biliarchiver.rest_api.bilivid import BiliVideo, VideoStatus from biliarchiver.version import BILI_ARCHIVER_VERSION @@ -19,19 +23,19 @@ from biliarchiver.version import BILI_ARCHIVER_VERSION @asynccontextmanager async def lifespan(app: FastAPI): global pending_queue, other_queue - pending_queue = BiliQueue() - other_queue = BiliQueue(maxsize=250) + pending_queue = BiliVideoQueue() + other_queue = BiliVideoQueue(maxsize=250) print("Loading queue...") load_queue() print("Queue loaded") - asyncio.create_task(scheduler()) + _video_scheduler = asyncio.create_task(video_scheduler()) yield print("Shutting down...") save_queue() print("Queue saved") -class BiliQueue(Queue): +class BiliVideoQueue(asyncio.Queue): def get_all(self) -> List[BiliVideo]: return list(self._queue) # type: ignore async def get(self) -> BiliVideo: @@ -49,8 +53,8 @@ class BiliQueue(Queue): ori_video.status = status await self.put(ori_video) -pending_queue: BiliQueue = None # type: ignore -other_queue: BiliQueue = None # type: ignore +pending_queue: BiliVideoQueue = None # type: ignore +other_queue: BiliVideoQueue = None # type: ignore app = FastAPI(lifespan=lifespan) @@ -102,8 +106,43 @@ async def delete(vid: str): return {"success": False, "vid": vid, "status": "not_found"} +async def source_action(fun, source_id: str, TRUNCATE=20): + try: + txt_path = await fun(source_id, truncate=TRUNCATE) + except Exception as e: + print(f"Failed to call {fun}: {e}") + raise HTTPException(status_code=500, detail=f"Failed to call {fun}") + if not isinstance(txt_path, Path): + raise HTTPException(status_code=500, detail="Failed to get path") -async def scheduler(): + bvids = read_bvids_from_txt(txt_path) + + return {"success": True, "bvids": bvids} + +@app.post("/get_bvids_by/{source_type}/{source_id}") +async def perform_source_action_from_req(source_type: str, source_id: str): + # make sure source_id is valid integer + if not source_id.isdecimal(): + raise HTTPException(status_code=400, detail="Invalid source_id") + + source_id = str(int(source_id)) + + fun_mapping = { + "up_videos": by_up_videos, + "favlist": by_favlist, + "series": by_series, + } + + if source_type not in fun_mapping: + raise HTTPException(status_code=400, detail="Invalid source_type") + + fun = fun_mapping[source_type] + + assert callable(fun) + + return await source_action(fun, source_id, TRUNCATE=int(9e99)) + +async def video_scheduler(): while True: print("Getting a video URL... If no video URL is printed, the queue is empty.") video = await pending_queue.get()