mirror of
https://github.com/jxxghp/MoviePilot.git
synced 2026-05-06 20:42:43 +08:00
10543eedd0
- 新增 /media/{mediaid}/stream 和 /title/stream 接口，支持基于 SSE 的渐进式搜索
- SearchChain 增加 async_search_by_title_stream、async_search_by_id_stream、async_process_stream、__async_search_all_sites_stream 方法
- 搜索结果按站点完成顺序实时推送，支持进度、候选、过滤、完成等阶段事件
- 优化参数解析与异常处理，提升大规模搜索体验
981 lines
41 KiB
Python
981 lines
41 KiB
Python
import asyncio
|
||
import random
|
||
import time
|
||
from concurrent.futures import ThreadPoolExecutor, as_completed
|
||
from datetime import datetime
|
||
from typing import AsyncIterator, Any, Dict, Tuple
|
||
from typing import List, Optional
|
||
|
||
from app.helper.sites import SitesHelper # noqa
|
||
from fastapi.concurrency import run_in_threadpool
|
||
|
||
from app.chain import ChainBase
|
||
from app.core.config import global_vars, settings
|
||
from app.core.context import Context
|
||
from app.core.context import MediaInfo, TorrentInfo
|
||
from app.core.event import eventmanager, Event
|
||
from app.core.metainfo import MetaInfo
|
||
from app.db.systemconfig_oper import SystemConfigOper
|
||
from app.helper.progress import ProgressHelper
|
||
from app.helper.torrent import TorrentHelper
|
||
from app.log import logger
|
||
from app.schemas import NotExistMediaInfo
|
||
from app.schemas.types import MediaType, ProgressKey, SystemConfigKey, EventType
|
||
|
||
|
||
class SearchChain(ChainBase):
|
||
"""
|
||
站点资源搜索处理链
|
||
"""
|
||
|
||
__result_temp_file = "__search_result__"
|
||
__ai_result_temp_file = "__ai_search_result__"
|
||
|
||
def search_by_id(self, tmdbid: Optional[int] = None, doubanid: Optional[str] = None,
|
||
mtype: MediaType = None, area: Optional[str] = "title", season: Optional[int] = None,
|
||
sites: List[int] = None, cache_local: bool = False) -> List[Context]:
|
||
"""
|
||
根据TMDBID/豆瓣ID搜索资源，精确匹配，不过滤本地存在的资源
|
||
:param tmdbid: TMDB ID
|
||
:param doubanid: 豆瓣 ID
|
||
:param mtype: 媒体，电影 or 电视剧
|
||
:param area: 搜索范围，title or imdbid
|
||
:param season: 季数
|
||
:param sites: 站点ID列表
|
||
:param cache_local: 是否缓存到本地
|
||
"""
|
||
mediainfo = self.recognize_media(tmdbid=tmdbid, doubanid=doubanid, mtype=mtype)
|
||
if not mediainfo:
|
||
logger.error(f'{tmdbid} 媒体信息识别失败！')
|
||
return []
|
||
no_exists = None
|
||
if season is not None:
|
||
no_exists = {
|
||
tmdbid or doubanid: {
|
||
season: NotExistMediaInfo(episodes=[])
|
||
}
|
||
}
|
||
results = self.process(mediainfo=mediainfo, sites=sites, area=area, no_exists=no_exists)
|
||
# 保存到本地文件
|
||
if cache_local:
|
||
self.save_cache(results, self.__result_temp_file)
|
||
return results
|
||
|
||
def search_by_title(self, title: str, page: Optional[int] = 0,
|
||
sites: List[int] = None, cache_local: Optional[bool] = False) -> List[Context]:
|
||
"""
|
||
根据标题搜索资源，不识别不过滤，直接返回站点内容
|
||
:param title: 标题，为空时返回所有站点首页内容
|
||
:param page: 页码
|
||
:param sites: 站点ID列表
|
||
:param cache_local: 是否缓存到本地
|
||
"""
|
||
if title:
|
||
logger.info(f'开始搜索资源，关键词：{title} ...')
|
||
else:
|
||
logger.info(f'开始浏览资源，站点：{sites} ...')
|
||
# 搜索
|
||
torrents = self.__search_all_sites(keyword=title, sites=sites, page=page) or []
|
||
if not torrents:
|
||
logger.warn(f'{title} 未搜索到资源')
|
||
return []
|
||
# 组装上下文
|
||
contexts = [Context(meta_info=MetaInfo(title=torrent.title, subtitle=torrent.description),
|
||
torrent_info=torrent) for torrent in torrents]
|
||
# 保存到本地文件
|
||
if cache_local:
|
||
self.save_cache(contexts, self.__result_temp_file)
|
||
return contexts
|
||
|
||
def last_search_results(self) -> Optional[List[Context]]:
|
||
"""
|
||
获取上次搜索结果
|
||
"""
|
||
return self.load_cache(self.__result_temp_file)
|
||
|
||
async def async_last_search_results(self) -> Optional[List[Context]]:
|
||
"""
|
||
异步获取上次搜索结果
|
||
"""
|
||
return await self.async_load_cache(self.__result_temp_file)
|
||
|
||
async def async_last_ai_results(self) -> Optional[List[Context]]:
|
||
"""
|
||
异步获取上次AI推荐结果
|
||
"""
|
||
return await self.async_load_cache(self.__ai_result_temp_file)
|
||
|
||
async def async_save_ai_results(self, results: List[Context]):
|
||
"""
|
||
异步保存AI推荐结果
|
||
"""
|
||
await self.async_save_cache(results, self.__ai_result_temp_file)
|
||
|
||
async def async_search_by_id(self, tmdbid: Optional[int] = None, doubanid: Optional[str] = None,
|
||
mtype: MediaType = None, area: Optional[str] = "title", season: Optional[int] = None,
|
||
sites: List[int] = None, cache_local: bool = False) -> List[Context]:
|
||
"""
|
||
根据TMDBID/豆瓣ID异步搜索资源，精确匹配，不过滤本地存在的资源
|
||
:param tmdbid: TMDB ID
|
||
:param doubanid: 豆瓣 ID
|
||
:param mtype: 媒体，电影 or 电视剧
|
||
:param area: 搜索范围，title or imdbid
|
||
:param season: 季数
|
||
:param sites: 站点ID列表
|
||
:param cache_local: 是否缓存到本地
|
||
"""
|
||
mediainfo = await self.async_recognize_media(tmdbid=tmdbid, doubanid=doubanid, mtype=mtype)
|
||
if not mediainfo:
|
||
logger.error(f'{tmdbid} 媒体信息识别失败！')
|
||
return []
|
||
no_exists = None
|
||
if season is not None:
|
||
no_exists = {
|
||
tmdbid or doubanid: {
|
||
season: NotExistMediaInfo(episodes=[])
|
||
}
|
||
}
|
||
results = await self.async_process(mediainfo=mediainfo, sites=sites, area=area, no_exists=no_exists)
|
||
# 保存到本地文件
|
||
if cache_local:
|
||
await self.async_save_cache(results, self.__result_temp_file)
|
||
return results
|
||
|
||
async def async_search_by_title(self, title: str, page: Optional[int] = 0,
|
||
sites: List[int] = None, cache_local: Optional[bool] = False) -> List[Context]:
|
||
"""
|
||
根据标题异步搜索资源，不识别不过滤，直接返回站点内容
|
||
:param title: 标题，为空时返回所有站点首页内容
|
||
:param page: 页码
|
||
:param sites: 站点ID列表
|
||
:param cache_local: 是否缓存到本地
|
||
"""
|
||
if title:
|
||
logger.info(f'开始搜索资源，关键词：{title} ...')
|
||
else:
|
||
logger.info(f'开始浏览资源，站点：{sites} ...')
|
||
# 搜索
|
||
torrents = await self.__async_search_all_sites(keyword=title, sites=sites, page=page) or []
|
||
if not torrents:
|
||
logger.warn(f'{title} 未搜索到资源')
|
||
return []
|
||
# 组装上下文
|
||
contexts = [Context(meta_info=MetaInfo(title=torrent.title, subtitle=torrent.description),
|
||
torrent_info=torrent) for torrent in torrents]
|
||
# 保存到本地文件
|
||
if cache_local:
|
||
await self.async_save_cache(contexts, self.__result_temp_file)
|
||
return contexts
|
||
|
||
async def async_search_by_title_stream(self, title: str, page: Optional[int] = 0,
|
||
sites: List[int] = None,
|
||
cache_local: Optional[bool] = False) -> AsyncIterator[dict]:
|
||
"""
|
||
根据标题渐进式搜索资源，不识别不过滤，按站点完成顺序返回结果
|
||
"""
|
||
if title:
|
||
logger.info(f'开始渐进式搜索资源，关键词：{title} ...')
|
||
else:
|
||
logger.info(f'开始渐进式浏览资源，站点：{sites} ...')
|
||
|
||
contexts: List[Context] = []
|
||
async for event in self.__async_search_all_sites_stream(keyword=title, sites=sites, page=page):
|
||
result = event.pop("items", []) or []
|
||
batch_contexts = [
|
||
Context(meta_info=MetaInfo(title=torrent.title, subtitle=torrent.description),
|
||
torrent_info=torrent)
|
||
for torrent in result
|
||
]
|
||
if batch_contexts:
|
||
contexts.extend(batch_contexts)
|
||
yield {
|
||
**event,
|
||
"type": "append",
|
||
"items": [context.to_dict() for context in batch_contexts],
|
||
"total_items": len(contexts)
|
||
}
|
||
|
||
if cache_local:
|
||
await self.async_save_cache(contexts, self.__result_temp_file)
|
||
|
||
if not contexts:
|
||
logger.warn(f'{title} 未搜索到资源')
|
||
yield {
|
||
"type": "done",
|
||
"text": f"搜索完成，共 {len(contexts)} 个资源",
|
||
"items": [context.to_dict() for context in contexts],
|
||
"total_items": len(contexts)
|
||
}
|
||
|
||
async def async_search_by_id_stream(self, tmdbid: Optional[int] = None, doubanid: Optional[str] = None,
|
||
mtype: MediaType = None, area: Optional[str] = "title",
|
||
season: Optional[int] = None, sites: List[int] = None,
|
||
cache_local: bool = False) -> AsyncIterator[dict]:
|
||
"""
|
||
根据TMDBID/豆瓣ID渐进式搜索资源，先返回站点原始候选，再返回过滤匹配后的最终结果
|
||
"""
|
||
mediainfo = await self.async_recognize_media(tmdbid=tmdbid, doubanid=doubanid, mtype=mtype)
|
||
if not mediainfo:
|
||
logger.error(f'{tmdbid} 媒体信息识别失败！')
|
||
yield {
|
||
"type": "error",
|
||
"success": False,
|
||
"message": "媒体信息识别失败"
|
||
}
|
||
return
|
||
|
||
no_exists = None
|
||
if season is not None:
|
||
no_exists = {
|
||
tmdbid or doubanid: {
|
||
season: NotExistMediaInfo(episodes=[])
|
||
}
|
||
}
|
||
|
||
contexts: List[Context] = []
|
||
async for event in self.async_process_stream(mediainfo=mediainfo, sites=sites, area=area, no_exists=no_exists):
|
||
if event.get("type") == "done":
|
||
contexts = event.get("contexts") or []
|
||
event = {
|
||
key: value
|
||
for key, value in event.items()
|
||
if key != "contexts"
|
||
}
|
||
yield event
|
||
|
||
if cache_local:
|
||
await self.async_save_cache(contexts, self.__result_temp_file)
|
||
|
||
@staticmethod
|
||
def __prepare_params(mediainfo: MediaInfo,
|
||
keyword: Optional[str] = None,
|
||
no_exists: Dict[int, Dict[int, NotExistMediaInfo]] = None
|
||
) -> Tuple[Dict[int, List[int]], List[str]]:
|
||
"""
|
||
准备搜索参数
|
||
"""
|
||
# 缺失的季集
|
||
mediakey = mediainfo.tmdb_id or mediainfo.douban_id
|
||
if no_exists and no_exists.get(mediakey):
|
||
# 过滤剧集
|
||
season_episodes = {sea: info.episodes
|
||
for sea, info in no_exists[mediakey].items()}
|
||
elif mediainfo.season is not None:
|
||
# 豆瓣只搜索当前季
|
||
season_episodes = {mediainfo.season: []}
|
||
else:
|
||
season_episodes = None
|
||
|
||
# 搜索关键词
|
||
if keyword:
|
||
keywords = [keyword]
|
||
else:
|
||
# 去重去空，但要保持顺序
|
||
keywords = list(dict.fromkeys([k for k in [mediainfo.title,
|
||
mediainfo.original_title,
|
||
mediainfo.en_title,
|
||
mediainfo.hk_title,
|
||
mediainfo.tw_title,
|
||
mediainfo.sg_title] if k]))
|
||
# 限制搜索关键词数量
|
||
if settings.MAX_SEARCH_NAME_LIMIT:
|
||
keywords = keywords[:settings.MAX_SEARCH_NAME_LIMIT]
|
||
|
||
return season_episodes, keywords
|
||
|
||
def __parse_result(self, torrents: List[TorrentInfo],
|
||
mediainfo: MediaInfo,
|
||
keyword: Optional[str] = None,
|
||
rule_groups: List[str] = None,
|
||
season_episodes: Dict[int, List[int]] = None,
|
||
custom_words: List[str] = None,
|
||
filter_params: Dict[str, str] = None) -> List[Context]:
|
||
"""
|
||
处理搜索结果
|
||
"""
|
||
|
||
def __do_filter(torrent_list: List[TorrentInfo]) -> List[TorrentInfo]:
|
||
"""
|
||
执行优先级过滤
|
||
"""
|
||
return self.filter_torrents(rule_groups=rule_groups,
|
||
torrent_list=torrent_list,
|
||
mediainfo=mediainfo) or []
|
||
|
||
if not torrents:
|
||
logger.warn(f'{keyword or mediainfo.title} 未搜索到资源')
|
||
return []
|
||
|
||
# 开始新进度
|
||
progress = ProgressHelper(ProgressKey.Search)
|
||
progress.start()
|
||
|
||
# 开始过滤
|
||
progress.update(value=0, text=f'开始过滤，总 {len(torrents)} 个资源，请稍候...')
|
||
# 匹配订阅附加参数
|
||
if filter_params:
|
||
logger.info(f'开始附加参数过滤，附加参数：{filter_params} ...')
|
||
torrents = [torrent for torrent in torrents if TorrentHelper().filter_torrent(torrent, filter_params)]
|
||
# 开始过滤规则过滤
|
||
if rule_groups is None:
|
||
# 取搜索过滤规则
|
||
rule_groups: List[str] = SystemConfigOper().get(SystemConfigKey.SearchFilterRuleGroups)
|
||
if rule_groups:
|
||
logger.info(f'开始过滤规则/剧集过滤，使用规则组：{rule_groups} ...')
|
||
torrents = __do_filter(torrents)
|
||
if not torrents:
|
||
logger.warn(f'{keyword or mediainfo.title} 没有符合过滤规则的资源')
|
||
return []
|
||
logger.info(f"过滤规则/剧集过滤完成，剩余 {len(torrents)} 个资源")
|
||
|
||
# 过滤完成
|
||
progress.update(value=50, text=f'过滤完成，剩余 {len(torrents)} 个资源')
|
||
|
||
# 总数
|
||
_total = len(torrents)
|
||
# 已处理数
|
||
_count = 0
|
||
|
||
# 开始匹配
|
||
_match_torrents = []
|
||
torrenthelper = TorrentHelper()
|
||
try:
|
||
# 英文标题应该在别名/原标题中，不需要再匹配
|
||
logger.info(f"开始匹配结果 标题：{mediainfo.title}，原标题：{mediainfo.original_title}，别名：{mediainfo.names}")
|
||
progress.update(value=51, text=f'开始匹配，总 {_total} 个资源 ...')
|
||
for torrent in torrents:
|
||
if global_vars.is_system_stopped:
|
||
break
|
||
_count += 1
|
||
progress.update(value=(_count / _total) * 96,
|
||
text=f'正在匹配 {torrent.site_name}，已完成 {_count} / {_total} ...')
|
||
if not torrent.title:
|
||
continue
|
||
|
||
# 识别元数据
|
||
torrent_meta = MetaInfo(title=torrent.title, subtitle=torrent.description,
|
||
custom_words=custom_words)
|
||
if torrent.title != torrent_meta.org_string:
|
||
logger.info(f"种子名称应用识别词后发生改变：{torrent.title} => {torrent_meta.org_string}")
|
||
# 季集数过滤
|
||
if season_episodes \
|
||
and not TorrentHelper.match_season_episodes(torrent=torrent,
|
||
meta=torrent_meta,
|
||
season_episodes=season_episodes):
|
||
continue
|
||
# 比对IMDBID
|
||
if torrent.imdbid \
|
||
and mediainfo.imdb_id \
|
||
and torrent.imdbid == mediainfo.imdb_id:
|
||
logger.info(f'{mediainfo.title} 通过IMDBID匹配到资源：{torrent.site_name} - {torrent.title}')
|
||
_match_torrents.append((torrent, torrent_meta))
|
||
continue
|
||
|
||
# 比对种子
|
||
if torrenthelper.match_torrent(mediainfo=mediainfo,
|
||
torrent_meta=torrent_meta,
|
||
torrent=torrent):
|
||
# 匹配成功
|
||
_match_torrents.append((torrent, torrent_meta))
|
||
continue
|
||
# 匹配完成
|
||
logger.info(f"匹配完成，共匹配到 {len(_match_torrents)} 个资源")
|
||
progress.update(value=97,
|
||
text=f'匹配完成，共匹配到 {len(_match_torrents)} 个资源')
|
||
|
||
# 去掉mediainfo中多余的数据
|
||
mediainfo.clear()
|
||
# 组装上下文
|
||
contexts = [Context(torrent_info=t[0],
|
||
media_info=mediainfo,
|
||
meta_info=t[1]) for t in _match_torrents]
|
||
finally:
|
||
torrents.clear()
|
||
del torrents
|
||
_match_torrents.clear()
|
||
del _match_torrents
|
||
|
||
# 排序
|
||
progress.update(value=99,
|
||
text=f'正在对 {len(contexts)} 个资源进行排序，请稍候...')
|
||
contexts = torrenthelper.sort_torrents(contexts)
|
||
|
||
# 结束进度
|
||
logger.info(f'搜索完成，共 {len(contexts)} 个资源')
|
||
progress.update(value=100,
|
||
text=f'搜索完成，共 {len(contexts)} 个资源')
|
||
progress.end()
|
||
|
||
# 去重后返回
|
||
return self.__remove_duplicate(contexts)
|
||
|
||
@staticmethod
|
||
def __remove_duplicate(_torrents: List[Context]) -> List[Context]:
|
||
"""
|
||
去除重复的种子
|
||
:param _torrents: 种子列表
|
||
:return: 去重后的种子列表
|
||
"""
|
||
return list({f"{t.torrent_info.site_name}_{t.torrent_info.title}_{t.torrent_info.description}": t
|
||
for t in _torrents}.values())
|
||
|
||
def process(self, mediainfo: MediaInfo,
|
||
keyword: Optional[str] = None,
|
||
no_exists: Dict[int, Dict[int, NotExistMediaInfo]] = None,
|
||
sites: List[int] = None,
|
||
rule_groups: List[str] = None,
|
||
area: Optional[str] = "title",
|
||
custom_words: List[str] = None,
|
||
filter_params: Dict[str, str] = None) -> List[Context]:
|
||
"""
|
||
根据媒体信息搜索种子资源，精确匹配，应用过滤规则，同时根据no_exists过滤本地已存在的资源
|
||
:param mediainfo: 媒体信息
|
||
:param keyword: 搜索关键词
|
||
:param no_exists: 缺失的媒体信息
|
||
:param sites: 站点ID列表，为空时搜索所有站点
|
||
:param rule_groups: 过滤规则组名称列表
|
||
:param area: 搜索范围，title or imdbid
|
||
:param custom_words: 自定义识别词列表
|
||
:param filter_params: 过滤参数
|
||
"""
|
||
|
||
# 豆瓣标题处理
|
||
if not mediainfo.tmdb_id:
|
||
meta = MetaInfo(title=mediainfo.title)
|
||
mediainfo.title = meta.name
|
||
mediainfo.season = meta.begin_season
|
||
logger.info(f'开始搜索资源，关键词：{keyword or mediainfo.title} ...')
|
||
|
||
# 补充媒体信息
|
||
if not mediainfo.names:
|
||
mediainfo: MediaInfo = self.recognize_media(mtype=mediainfo.type,
|
||
tmdbid=mediainfo.tmdb_id,
|
||
doubanid=mediainfo.douban_id)
|
||
if not mediainfo:
|
||
logger.error(f'媒体信息识别失败！')
|
||
return []
|
||
|
||
# 准备搜索参数
|
||
season_episodes, keywords = self.__prepare_params(
|
||
mediainfo=mediainfo,
|
||
keyword=keyword,
|
||
no_exists=no_exists
|
||
)
|
||
|
||
# 站点搜索结果
|
||
torrents: List[TorrentInfo] = []
|
||
# 站点搜索次数
|
||
search_count = 0
|
||
|
||
# 多关键字执行搜索
|
||
for search_word in keywords:
|
||
# 强制休眠 1-10 秒
|
||
if search_count > 0:
|
||
logger.info(f"已搜索 {search_count} 次，强制休眠 1-10 秒 ...")
|
||
time.sleep(random.randint(1, 10))
|
||
|
||
# 搜索站点
|
||
results = self.__search_all_sites(
|
||
mediainfo=mediainfo,
|
||
keyword=search_word,
|
||
sites=sites,
|
||
area=area
|
||
) or []
|
||
# 合并结果
|
||
|
||
search_count += 1
|
||
torrents.extend(results)
|
||
|
||
# 有结果则停止
|
||
if not settings.SEARCH_MULTIPLE_NAME and torrents:
|
||
logger.info(f"共搜索到 {len(torrents)} 个资源，停止搜索")
|
||
break
|
||
|
||
# 处理结果
|
||
return self.__parse_result(
|
||
torrents=torrents,
|
||
mediainfo=mediainfo,
|
||
keyword=keyword,
|
||
rule_groups=rule_groups,
|
||
season_episodes=season_episodes,
|
||
custom_words=custom_words,
|
||
filter_params=filter_params
|
||
)
|
||
|
||
async def async_process(self, mediainfo: MediaInfo,
|
||
keyword: Optional[str] = None,
|
||
no_exists: Dict[int, Dict[int, NotExistMediaInfo]] = None,
|
||
sites: List[int] = None,
|
||
rule_groups: List[str] = None,
|
||
area: Optional[str] = "title",
|
||
custom_words: List[str] = None,
|
||
filter_params: Dict[str, str] = None) -> List[Context]:
|
||
"""
|
||
根据媒体信息异步搜索种子资源，精确匹配，应用过滤规则，同时根据no_exists过滤本地已存在的资源
|
||
:param mediainfo: 媒体信息
|
||
:param keyword: 搜索关键词
|
||
:param no_exists: 缺失的媒体信息
|
||
:param sites: 站点ID列表，为空时搜索所有站点
|
||
:param rule_groups: 过滤规则组名称列表
|
||
:param area: 搜索范围，title or imdbid
|
||
:param custom_words: 自定义识别词列表
|
||
:param filter_params: 过滤参数
|
||
"""
|
||
|
||
# 豆瓣标题处理
|
||
if not mediainfo.tmdb_id:
|
||
meta = MetaInfo(title=mediainfo.title)
|
||
mediainfo.title = meta.name
|
||
mediainfo.season = meta.begin_season
|
||
logger.info(f'开始搜索资源，关键词：{keyword or mediainfo.title} ...')
|
||
|
||
# 补充媒体信息
|
||
if not mediainfo.names:
|
||
mediainfo: MediaInfo = await self.async_recognize_media(mtype=mediainfo.type,
|
||
tmdbid=mediainfo.tmdb_id,
|
||
doubanid=mediainfo.douban_id)
|
||
if not mediainfo:
|
||
logger.error(f'媒体信息识别失败！')
|
||
return []
|
||
|
||
# 准备搜索参数
|
||
season_episodes, keywords = self.__prepare_params(
|
||
mediainfo=mediainfo,
|
||
keyword=keyword,
|
||
no_exists=no_exists
|
||
)
|
||
|
||
# 站点搜索结果
|
||
torrents: List[TorrentInfo] = []
|
||
# 站点搜索次数
|
||
search_count = 0
|
||
|
||
# 多关键字执行搜索
|
||
for search_word in keywords:
|
||
# 强制休眠 1-10 秒
|
||
if search_count > 0:
|
||
logger.info(f"已搜索 {search_count} 次，强制休眠 1-10 秒 ...")
|
||
await asyncio.sleep(random.randint(1, 10))
|
||
# 搜索站点
|
||
torrents.extend(
|
||
await self.__async_search_all_sites(
|
||
mediainfo=mediainfo,
|
||
keyword=search_word,
|
||
sites=sites,
|
||
area=area
|
||
) or []
|
||
)
|
||
search_count += 1
|
||
# 有结果则停止
|
||
if torrents:
|
||
logger.info(f"共搜索到 {len(torrents)} 个资源，停止搜索")
|
||
break
|
||
|
||
# 处理结果
|
||
return await run_in_threadpool(self.__parse_result,
|
||
torrents=torrents,
|
||
mediainfo=mediainfo,
|
||
keyword=keyword,
|
||
rule_groups=rule_groups,
|
||
season_episodes=season_episodes,
|
||
custom_words=custom_words,
|
||
filter_params=filter_params
|
||
)
|
||
|
||
async def async_process_stream(self, mediainfo: MediaInfo,
|
||
keyword: Optional[str] = None,
|
||
no_exists: Dict[int, Dict[int, NotExistMediaInfo]] = None,
|
||
sites: List[int] = None,
|
||
rule_groups: List[str] = None,
|
||
area: Optional[str] = "title",
|
||
custom_words: List[str] = None,
|
||
filter_params: Dict[str, str] = None) -> AsyncIterator[dict]:
|
||
"""
|
||
根据媒体信息渐进式搜索种子资源，先返回站点候选，再返回过滤匹配后的最终结果
|
||
"""
|
||
|
||
# 豆瓣标题处理
|
||
if not mediainfo.tmdb_id:
|
||
meta = MetaInfo(title=mediainfo.title)
|
||
mediainfo.title = meta.name
|
||
mediainfo.season = meta.begin_season
|
||
logger.info(f'开始渐进式搜索资源，关键词：{keyword or mediainfo.title} ...')
|
||
|
||
# 补充媒体信息
|
||
if not mediainfo.names:
|
||
mediainfo = await self.async_recognize_media(mtype=mediainfo.type,
|
||
tmdbid=mediainfo.tmdb_id,
|
||
doubanid=mediainfo.douban_id)
|
||
if not mediainfo:
|
||
logger.error(f'媒体信息识别失败！')
|
||
yield {
|
||
"type": "error",
|
||
"success": False,
|
||
"message": "媒体信息识别失败"
|
||
}
|
||
return
|
||
|
||
# 准备搜索参数
|
||
season_episodes, keywords = self.__prepare_params(
|
||
mediainfo=mediainfo,
|
||
keyword=keyword,
|
||
no_exists=no_exists
|
||
)
|
||
|
||
torrents: List[TorrentInfo] = []
|
||
candidate_contexts: List[Context] = []
|
||
search_count = 0
|
||
|
||
for search_word in keywords:
|
||
if search_count > 0:
|
||
logger.info(f"已搜索 {search_count} 次，强制休眠 1-10 秒 ...")
|
||
await asyncio.sleep(random.randint(1, 10))
|
||
|
||
async for event in self.__async_search_all_sites_stream(
|
||
mediainfo=mediainfo,
|
||
keyword=search_word,
|
||
sites=sites,
|
||
area=area):
|
||
result = event.pop("items", []) or []
|
||
torrents.extend(result)
|
||
batch_contexts = [
|
||
Context(meta_info=MetaInfo(title=torrent.title, subtitle=torrent.description),
|
||
media_info=mediainfo,
|
||
torrent_info=torrent)
|
||
for torrent in result
|
||
]
|
||
candidate_contexts.extend(batch_contexts)
|
||
yield {
|
||
**event,
|
||
"type": "append",
|
||
"stage": "searching",
|
||
"items": [context.to_dict() for context in batch_contexts],
|
||
"total_items": len(candidate_contexts)
|
||
}
|
||
|
||
search_count += 1
|
||
if torrents:
|
||
logger.info(f"共搜索到 {len(torrents)} 个资源，停止搜索")
|
||
break
|
||
|
||
yield {
|
||
"type": "progress",
|
||
"stage": "filtering",
|
||
"value": 98,
|
||
"text": f"正在过滤匹配 {len(torrents)} 个候选资源 ..."
|
||
}
|
||
|
||
contexts = await run_in_threadpool(self.__parse_result,
|
||
torrents=torrents,
|
||
mediainfo=mediainfo,
|
||
keyword=keyword,
|
||
rule_groups=rule_groups,
|
||
season_episodes=season_episodes,
|
||
custom_words=custom_words,
|
||
filter_params=filter_params)
|
||
final_items = [context.to_dict() for context in contexts]
|
||
yield {
|
||
"type": "replace",
|
||
"stage": "filtered",
|
||
"value": 100,
|
||
"text": f"过滤匹配完成，共 {len(contexts)} 个资源",
|
||
"items": final_items,
|
||
"total_items": len(contexts)
|
||
}
|
||
yield {
|
||
"type": "done",
|
||
"stage": "done",
|
||
"text": f"搜索完成，共 {len(contexts)} 个资源",
|
||
"items": final_items,
|
||
"total_items": len(contexts),
|
||
"contexts": contexts
|
||
}
|
||
|
||
def __search_all_sites(self, keyword: str,
|
||
mediainfo: Optional[MediaInfo] = None,
|
||
sites: List[int] = None,
|
||
page: Optional[int] = 0,
|
||
area: Optional[str] = "title") -> Optional[List[TorrentInfo]]:
|
||
"""
|
||
多线程搜索多个站点
|
||
:param mediainfo: 识别的媒体信息
|
||
:param keyword: 搜索关键词
|
||
:param sites: 指定站点ID列表，如有则只搜索指定站点，否则搜索所有站点
|
||
:param page: 搜索页码
|
||
:param area: 搜索区域 title or imdbid
|
||
:reutrn: 资源列表
|
||
"""
|
||
# 未开启的站点不搜索
|
||
indexer_sites = []
|
||
|
||
# 配置的索引站点
|
||
if not sites:
|
||
sites = SystemConfigOper().get(SystemConfigKey.IndexerSites) or []
|
||
|
||
for indexer in SitesHelper().get_indexers():
|
||
# 检查站点索引开关
|
||
if not sites or indexer.get("id") in sites:
|
||
indexer_sites.append(indexer)
|
||
if not indexer_sites:
|
||
logger.warn('未开启任何有效站点，无法搜索资源')
|
||
return []
|
||
|
||
# 开始进度
|
||
progress = ProgressHelper(ProgressKey.Search)
|
||
progress.start()
|
||
# 开始计时
|
||
start_time = datetime.now()
|
||
# 总数
|
||
total_num = len(indexer_sites)
|
||
# 完成数
|
||
finish_count = 0
|
||
# 更新进度
|
||
progress.update(value=0,
|
||
text=f"开始搜索，共 {total_num} 个站点 ...")
|
||
# 结果集
|
||
results = []
|
||
# 多线程
|
||
with ThreadPoolExecutor(max_workers=len(indexer_sites)) as executor:
|
||
all_task = []
|
||
for site in indexer_sites:
|
||
if area == "imdbid":
|
||
# 搜索IMDBID
|
||
task = executor.submit(self.search_torrents, site=site,
|
||
keyword=mediainfo.imdb_id if mediainfo else None,
|
||
mtype=mediainfo.type if mediainfo else None,
|
||
page=page)
|
||
else:
|
||
# 搜索标题
|
||
task = executor.submit(self.search_torrents, site=site,
|
||
keyword=keyword,
|
||
mtype=mediainfo.type if mediainfo else None,
|
||
page=page)
|
||
all_task.append(task)
|
||
for future in as_completed(all_task):
|
||
if global_vars.is_system_stopped:
|
||
break
|
||
finish_count += 1
|
||
result = future.result()
|
||
if result:
|
||
results.extend(result)
|
||
logger.info(f"站点搜索进度：{finish_count} / {total_num}")
|
||
progress.update(value=finish_count / total_num * 100,
|
||
text=f"正在搜索{keyword or ''}，已完成 {finish_count} / {total_num} 个站点 ...")
|
||
# 计算耗时
|
||
end_time = datetime.now()
|
||
# 更新进度
|
||
progress.update(value=100,
|
||
text=f"站点搜索完成，有效资源数：{len(results)}，总耗时 {(end_time - start_time).seconds} 秒")
|
||
logger.info(f"站点搜索完成，有效资源数：{len(results)}，总耗时 {(end_time - start_time).seconds} 秒")
|
||
# 结束进度
|
||
progress.end()
|
||
|
||
# 返回
|
||
return results
|
||
|
||
async def __async_search_all_sites(self, keyword: str,
|
||
mediainfo: Optional[MediaInfo] = None,
|
||
sites: List[int] = None,
|
||
page: Optional[int] = 0,
|
||
area: Optional[str] = "title") -> Optional[List[TorrentInfo]]:
|
||
"""
|
||
异步搜索多个站点
|
||
:param mediainfo: 识别的媒体信息
|
||
:param keyword: 搜索关键词
|
||
:param sites: 指定站点ID列表，如有则只搜索指定站点，否则搜索所有站点
|
||
:param page: 搜索页码
|
||
:param area: 搜索区域 title or imdbid
|
||
:reutrn: 资源列表
|
||
"""
|
||
# 未开启的站点不搜索
|
||
indexer_sites = []
|
||
|
||
# 配置的索引站点
|
||
if not sites:
|
||
sites = SystemConfigOper().get(SystemConfigKey.IndexerSites) or []
|
||
|
||
for indexer in await SitesHelper().async_get_indexers():
|
||
# 检查站点索引开关
|
||
if not sites or indexer.get("id") in sites:
|
||
indexer_sites.append(indexer)
|
||
if not indexer_sites:
|
||
logger.warn('未开启任何有效站点，无法搜索资源')
|
||
return []
|
||
|
||
# 开始进度
|
||
progress = ProgressHelper(ProgressKey.Search)
|
||
progress.start()
|
||
# 开始计时
|
||
start_time = datetime.now()
|
||
# 总数
|
||
total_num = len(indexer_sites)
|
||
# 完成数
|
||
finish_count = 0
|
||
# 更新进度
|
||
progress.update(value=0,
|
||
text=f"开始搜索，共 {total_num} 个站点 ...")
|
||
# 结果集
|
||
results = []
|
||
|
||
# 创建异步任务列表
|
||
tasks = []
|
||
for site in indexer_sites:
|
||
if area == "imdbid":
|
||
# 搜索IMDBID
|
||
task = self.async_search_torrents(site=site,
|
||
keyword=mediainfo.imdb_id if mediainfo else None,
|
||
mtype=mediainfo.type if mediainfo else None,
|
||
page=page)
|
||
else:
|
||
# 搜索标题
|
||
task = self.async_search_torrents(site=site,
|
||
keyword=keyword,
|
||
mtype=mediainfo.type if mediainfo else None,
|
||
page=page)
|
||
tasks.append(task)
|
||
|
||
# 使用asyncio.as_completed来处理并发任务
|
||
for future in asyncio.as_completed(tasks):
|
||
if global_vars.is_system_stopped:
|
||
break
|
||
finish_count += 1
|
||
result = await future
|
||
if result:
|
||
results.extend(result)
|
||
logger.info(f"站点搜索进度：{finish_count} / {total_num}")
|
||
progress.update(value=finish_count / total_num * 100,
|
||
text=f"正在搜索{keyword or ''}，已完成 {finish_count} / {total_num} 个站点 ...")
|
||
|
||
# 计算耗时
|
||
end_time = datetime.now()
|
||
# 更新进度
|
||
progress.update(value=100,
|
||
text=f"站点搜索完成，有效资源数：{len(results)}，总耗时 {(end_time - start_time).seconds} 秒")
|
||
logger.info(f"站点搜索完成，有效资源数：{len(results)}，总耗时 {(end_time - start_time).seconds} 秒")
|
||
# 结束进度
|
||
progress.end()
|
||
|
||
# 返回
|
||
return results
|
||
|
||
async def __async_search_all_sites_stream(self, keyword: str,
|
||
mediainfo: Optional[MediaInfo] = None,
|
||
sites: List[int] = None,
|
||
page: Optional[int] = 0,
|
||
area: Optional[str] = "title") -> AsyncIterator[Dict[str, Any]]:
|
||
"""
|
||
异步搜索多个站点，按站点完成顺序渐进式返回结果
|
||
:param mediainfo: 识别的媒体信息
|
||
:param keyword: 搜索关键词
|
||
:param sites: 指定站点ID列表，如有则只搜索指定站点，否则搜索所有站点
|
||
:param page: 搜索页码
|
||
:param area: 搜索区域 title or imdbid
|
||
"""
|
||
indexer_sites = []
|
||
|
||
if not sites:
|
||
sites = SystemConfigOper().get(SystemConfigKey.IndexerSites) or []
|
||
|
||
for indexer in await SitesHelper().async_get_indexers():
|
||
if not sites or indexer.get("id") in sites:
|
||
indexer_sites.append(indexer)
|
||
if not indexer_sites:
|
||
logger.warn('未开启任何有效站点，无法搜索资源')
|
||
yield {
|
||
"type": "done",
|
||
"stage": "searching",
|
||
"value": 100,
|
||
"text": "未开启任何有效站点，无法搜索资源",
|
||
"items": [],
|
||
"finished": 0,
|
||
"total": 0
|
||
}
|
||
return
|
||
|
||
progress = ProgressHelper(ProgressKey.Search)
|
||
progress.start()
|
||
start_time = datetime.now()
|
||
total_num = len(indexer_sites)
|
||
finish_count = 0
|
||
progress.update(value=0,
|
||
text=f"开始搜索，共 {total_num} 个站点 ...")
|
||
yield {
|
||
"type": "progress",
|
||
"stage": "searching",
|
||
"value": 0,
|
||
"text": f"开始搜索，共 {total_num} 个站点 ...",
|
||
"items": [],
|
||
"finished": 0,
|
||
"total": total_num
|
||
}
|
||
|
||
async def search_site(site: dict) -> Tuple[dict, List[TorrentInfo]]:
|
||
if area == "imdbid":
|
||
result = await self.async_search_torrents(site=site,
|
||
keyword=mediainfo.imdb_id if mediainfo else None,
|
||
mtype=mediainfo.type if mediainfo else None,
|
||
page=page)
|
||
else:
|
||
result = await self.async_search_torrents(site=site,
|
||
keyword=keyword,
|
||
mtype=mediainfo.type if mediainfo else None,
|
||
page=page)
|
||
return site, result or []
|
||
|
||
tasks = [asyncio.create_task(search_site(site)) for site in indexer_sites]
|
||
results_count = 0
|
||
try:
|
||
for future in asyncio.as_completed(tasks):
|
||
if global_vars.is_system_stopped:
|
||
break
|
||
finish_count += 1
|
||
site, result = await future
|
||
results_count += len(result)
|
||
logger.info(f"站点搜索进度：{finish_count} / {total_num}")
|
||
progress_value = finish_count / total_num * 100
|
||
progress_text = f"正在搜索{keyword or ''}，已完成 {finish_count} / {total_num} 个站点 ..."
|
||
progress.update(value=progress_value, text=progress_text)
|
||
yield {
|
||
"type": "append",
|
||
"stage": "searching",
|
||
"value": progress_value,
|
||
"text": progress_text,
|
||
"items": result,
|
||
"site": site.get("name"),
|
||
"site_id": site.get("id"),
|
||
"finished": finish_count,
|
||
"total": total_num,
|
||
"total_items": results_count
|
||
}
|
||
finally:
|
||
for task in tasks:
|
||
if not task.done():
|
||
task.cancel()
|
||
|
||
end_time = datetime.now()
|
||
progress.update(value=100,
|
||
text=f"站点搜索完成，有效资源数：{results_count}，总耗时 {(end_time - start_time).seconds} 秒")
|
||
logger.info(f"站点搜索完成，有效资源数：{results_count}，总耗时 {(end_time - start_time).seconds} 秒")
|
||
progress.end()
|
||
|
||
@eventmanager.register(EventType.SiteDeleted)
|
||
def remove_site(self, event: Event):
|
||
"""
|
||
从搜索站点中移除与已删除站点相关的设置
|
||
"""
|
||
if not event:
|
||
return
|
||
event_data = event.event_data or {}
|
||
site_id = event_data.get("site_id")
|
||
if not site_id:
|
||
return
|
||
if site_id == "*":
|
||
# 清空搜索站点
|
||
SystemConfigOper().set(SystemConfigKey.IndexerSites, [])
|
||
return
|
||
# 从选中的rss站点中移除
|
||
selected_sites = SystemConfigOper().get(SystemConfigKey.IndexerSites) or []
|
||
if site_id in selected_sites:
|
||
selected_sites.remove(site_id)
|
||
SystemConfigOper().set(SystemConfigKey.IndexerSites, selected_sites)
|