Merge pull request #18 from saveweb/add-from-source

feat: Add API endpoints for archiving from user, favlist, and collection

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,

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

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()