diff --git a/backend/src/module/api/bangumi.py b/backend/src/module/api/bangumi.py index a96c22c2..c912f978 100644 --- a/backend/src/module/api/bangumi.py +++ b/backend/src/module/api/bangumi.py @@ -75,6 +75,13 @@ async def enable_rule(bangumi_id: str): return u_response(resp) +@router.get(path="/refresh/poster/all", response_model=APIResponse, dependencies=[Depends(get_current_user)]) +async def refresh_poster(): + with TorrentManager() as manager: + resp = manager.refresh_poster() + return u_response(resp) + + @router.get("/reset/all", response_model=APIResponse, dependencies=[Depends(get_current_user)]) async def reset_all(): with TorrentManager() as manager: diff --git a/backend/src/module/manager/collector.py b/backend/src/module/manager/collector.py index f6f1dab8..927f5919 100644 --- a/backend/src/module/manager/collector.py +++ b/backend/src/module/manager/collector.py @@ -9,13 +9,13 @@ logger = logging.getLogger(__name__) class SeasonCollector(DownloadClient): - def collect_season(self, bangumi: Bangumi, link: str = None): + def collect_season(self, bangumi: Bangumi, link: str = None, limit: int = 50): logger.info( f"Start collecting {bangumi.official_title} Season {bangumi.season}..." ) with SearchTorrent() as st: if not link: - torrents = st.search_season(bangumi) + torrents = st.search_season(bangumi, limit=limit) else: torrents = st.get_torrents(link, bangumi.filter.replace(",", "|")) if self.add_torrent(torrents, bangumi): @@ -23,6 +23,7 @@ class SeasonCollector(DownloadClient): bangumi.eps_collect = True with RSSEngine() as engine: engine.bangumi.update(bangumi) + engine.torrent.add_all(torrents) return ResponseModel( status=True, status_code=200, @@ -57,7 +58,7 @@ def eps_complete(): logger.info("Start collecting full season...") for data in datas: if not data.eps_collect: - with SeasonCollector() as sc: - sc.collect_season(data) + with SeasonCollector() as collector: + collector.collect_season(data) data.eps_collect = True engine.bangumi.update_all(datas) diff --git a/backend/src/module/manager/torrent.py b/backend/src/module/manager/torrent.py index 9a559fbc..23e1d2fe 100644 --- a/backend/src/module/manager/torrent.py +++ b/backend/src/module/manager/torrent.py @@ -1,10 +1,10 @@ import logging -from fastapi.responses import JSONResponse from module.database import Database from module.downloader import DownloadClient from module.models import Bangumi, BangumiUpdate, ResponseModel +from module.parser import TitleParser logger = logging.getLogger(__name__) @@ -133,6 +133,19 @@ class TorrentManager(Database): msg_zh=f"无法找到 id {bangumi_id} 的数据", ) + def refresh_poster(self): + bangumis = self.bangumi.search_all() + for bangumi in bangumis: + if not bangumi.poster_link: + TitleParser().tmdb_poster_parser(bangumi) + self.bangumi.update_all(bangumis) + return ResponseModel( + status_code=200, + status=True, + msg_en=f"Refresh poster link successfully.", + msg_zh=f"刷新海报链接成功。", + ) + def search_all_bangumi(self): datas = self.bangumi.search_all() diff --git a/backend/src/module/parser/title_parser.py b/backend/src/module/parser/title_parser.py index 89f0f2b2..1730ee31 100644 --- a/backend/src/module/parser/title_parser.py +++ b/backend/src/module/parser/title_parser.py @@ -36,6 +36,16 @@ class TitleParser: logger.warning("Please change bangumi info manually.") return title, season, None, None + @staticmethod + def tmdb_poster_parser(bangumi: Bangumi): + tmdb_info = tmdb_parser(bangumi.official_title, settings.rss_parser.language) + if tmdb_info: + logger.debug(f"TMDB Matched, official title is {tmdb_info.title}") + bangumi.poster_link = tmdb_info.poster_link + else: + logger.warning(f"Cannot match {bangumi.official_title} in TMDB. Use raw title instead.") + logger.warning("Please change bangumi info manually.") + @staticmethod def raw_parser(raw: str) -> Bangumi | None: language = settings.rss_parser.language diff --git a/backend/src/module/searcher/searcher.py b/backend/src/module/searcher/searcher.py index e38b7402..2dacbe40 100644 --- a/backend/src/module/searcher/searcher.py +++ b/backend/src/module/searcher/searcher.py @@ -44,7 +44,7 @@ class SearchTorrent(RequestContent, RSSAnalyser): url = search_url(site, keywords) return url - def search_season(self, data: Bangumi, site: str = "mikan") -> list[Torrent]: + def search_season(self, data: Bangumi, site: str = "mikan", limit: int = 5) -> list[Torrent]: rss_item = self.special_url(data, site) - torrents = self.search_torrents(rss_item) + torrents = self.search_torrents(rss_item, limit=limit) return [torrent for torrent in torrents if data.title_raw in torrent.name] \ No newline at end of file