diff --git a/app/helper/message.py b/app/helper/message.py index 922b43bb..c1c2cc09 100644 --- a/app/helper/message.py +++ b/app/helper/message.py @@ -28,11 +28,16 @@ from app.utils.string import StringUtils class TemplateContextBuilder: """ - 模板上下文构建器 - """ + 模板上下文构建器。 - def __init__(self): - self._context = {} + 无状态实现：所有 ``_add_*`` 方法均为静态方法，接受并就地修改调用方提供的 + ``context`` 字典。``build`` 每次调用都基于一份新的本地字典装填后返回， + 实例自身不持有任何中间状态——可以被多线程共享调用而不会产生互相串味的 + ``rename_dict``，配合 ``settings.TRANSFER_THREADS > 1`` 的并发整理场景安全。 + + 保留为类（而非自由函数）是为了向后兼容现有调用方式 + （``TemplateHelper().builder.build(...)``）。 + """ def build( self, @@ -46,55 +51,64 @@ class TemplateContextBuilder: **kwargs ) -> Dict[str, Any]: """ - :param meta: 媒体信息 - :param mediainfo: 媒体信息 + 构建一次性渲染上下文字典。 + + 每次调用都新建本地 ``context`` 字典，依次填充各业务来源后返回过滤掉 + None 值的副本，调用之间互不影响。 + + :param meta: 媒体元数据 + :param mediainfo: 识别的媒体信息 :param torrentinfo: 种子信息 - :param transferinfo: 传输信息 + :param transferinfo: 整理结果信息 :param file_extension: 文件扩展名 - :param episodes_info: 剧集信息 - :param include_raw_objects: 是否包含原始对象 + :param episodes_info: 当前季的全部集信息 + :param include_raw_objects: 是否在 dict 里附带原始对象引用（``__meta__`` 等） :return: 渲染上下文字典 """ - self._context.clear() - self._add_episode_details(meta, episodes_info) - self._add_media_info(mediainfo) - self._add_transfer_info(transferinfo) - self._add_torrent_info(torrentinfo) - self._add_file_info(file_extension) + context: Dict[str, Any] = {} + self._add_episode_details(context, meta, episodes_info) + self._add_media_info(context, mediainfo) + self._add_transfer_info(context, transferinfo) + self._add_torrent_info(context, torrentinfo) + self._add_file_info(context, file_extension) if kwargs: - self._context.update(kwargs) + context.update(kwargs) if include_raw_objects: - self._add_raw_objects(meta, mediainfo, torrentinfo, transferinfo, episodes_info) + self._add_raw_objects(context, meta, mediainfo, torrentinfo, transferinfo, episodes_info) # 移除空值 - return {k: v for k, v in self._context.items() if v is not None} + return {k: v for k, v in context.items() if v is not None} - def _add_media_info(self, mediainfo: MediaInfo): + @classmethod + def _add_media_info(cls, context: Dict[str, Any], mediainfo: Optional[MediaInfo]) -> None: """ - 增加媒体信息 + 将 MediaInfo 中的标题、季年份、海报等业务字段就地写入 ``context``。 + + 会读取 ``context`` 中由 ``_add_episode_details`` 先填好的 ``season`` / + ``year`` / ``title_year`` 占位，保证电视剧场景下季/年优先沿用 meta 解析值。 """ if not mediainfo: return season_fmt = f"S{mediainfo.season:02d}" if mediainfo.season is not None else None base_info = { # 标题 - "title": self.__convert_invalid_characters(mediainfo.title), + "title": cls.__convert_invalid_characters(mediainfo.title), # 英文标题 - "en_title": self.__convert_invalid_characters(mediainfo.en_title), + "en_title": cls.__convert_invalid_characters(mediainfo.en_title), # 原语种标题 - "original_title": self.__convert_invalid_characters(mediainfo.original_title), + "original_title": cls.__convert_invalid_characters(mediainfo.original_title), # 季号 - "season": self._context.get("season") or mediainfo.season, + "season": context.get("season") or mediainfo.season, # Sxx - "season_fmt": self._context.get("season_fmt") or season_fmt, + "season_fmt": context.get("season_fmt") or season_fmt, # 年份 - "year": mediainfo.year or self._context.get("year"), + "year": mediainfo.year or context.get("year"), # 媒体标题 + 年份 - "title_year": mediainfo.title_year or self._context.get("title_year"), + "title_year": mediainfo.title_year or context.get("title_year"), } - _meta_season = self._context.get("season") + _meta_season = context.get("season") media_info = { # 类型 "type": mediainfo.type.value, @@ -121,11 +135,18 @@ class TemplateContextBuilder: # 豆瓣ID "doubanid": mediainfo.douban_id, } - self._context.update({**base_info, **media_info}) + context.update({**base_info, **media_info}) - def _add_episode_details(self, meta: Optional[MetaBase], episodes: Optional[List[TmdbEpisode]]): + @classmethod + def _add_episode_details( + cls, + context: Dict[str, Any], + meta: Optional[MetaBase], + episodes: Optional[List[TmdbEpisode]], + ) -> None: """ - 添加剧集详细信息 + 将 meta 解析得到的剧集级信息、技术字段写入 ``context``，并尝试匹配 + TMDB 集详情填入 ``episode_title`` / ``episode_date``。 """ if not meta: return @@ -135,7 +156,7 @@ class TemplateContextBuilder: for episode in episodes: if episode.episode_number == meta.begin_episode: episode_data.update({ - "episode_title": self.__convert_invalid_characters(episode.name), + "episode_title": cls.__convert_invalid_characters(episode.name), "episode_date": episode.air_date if episode.air_date else None }) break @@ -150,7 +171,7 @@ class TemplateContextBuilder: # 年份 "year": meta.year, # 名字 + 年份 - "title_year": self._context.get("title_year") or "%s (%s)" % ( + "title_year": context.get("title_year") or "%s (%s)" % ( meta.name, meta.year) if meta.year else meta.name, # 季号 "season": meta.season_seq, @@ -188,11 +209,16 @@ class TemplateContextBuilder: # 流媒体平台 "webSource": meta.web_source, } - self._context.update({**meta_info, **tech_metadata, **episode_data}) + context.update({**meta_info, **tech_metadata, **episode_data}) - def _add_torrent_info(self, torrentinfo: Optional[TorrentInfo]): + @staticmethod + def _add_torrent_info(context: Dict[str, Any], torrentinfo: Optional[TorrentInfo]) -> None: """ - 添加种子信息 + 将种子信息写入 ``context``，描述字段会去除 HTML 标签。 + + 副作用提醒：当 ``torrentinfo.description`` 包含 HTML 时，会就地清洗 + 原对象的 description 字段——保留原始行为，避免破坏现有调用方对清洗后 + 描述的依赖。 """ if not torrentinfo: return @@ -231,25 +257,27 @@ class TemplateContextBuilder: # 种子大小 "size": size, } - self._context.update(torrent_info) + context.update(torrent_info) - def _add_transfer_info(self, transferinfo: Optional[TransferInfo]) -> Optional[Dict]: + @staticmethod + def _add_transfer_info(context: Dict[str, Any], transferinfo: Optional[TransferInfo]) -> None: """ - 添加文件转移上下文 + 将整理结果（类型、文件数、总大小、错误信息）写入 ``context``。 """ if not transferinfo: - return None + return ctx = { "transfer_type": transferinfo.transfer_type, "file_count": transferinfo.file_count, "total_size": StringUtils.str_filesize(transferinfo.total_size), "err_msg": transferinfo.message, } - return self._context.update(ctx) + context.update(ctx) - def _add_file_info(self, file_extension: Optional[str]): + @staticmethod + def _add_file_info(context: Dict[str, Any], file_extension: Optional[str]) -> None: """ - 添加文件信息 + 将文件扩展名写入 ``context.fileExt``。 """ if not file_extension: return @@ -257,18 +285,21 @@ class TemplateContextBuilder: # 文件后缀 "fileExt": file_extension, } - self._context.update(file_info) + context.update(file_info) + @staticmethod def _add_raw_objects( - self, + context: Dict[str, Any], meta: Optional[MetaBase], mediainfo: Optional[MediaInfo], torrentinfo: Optional[TorrentInfo], transferinfo: Optional[TransferInfo], episodes_info: Optional[List[TmdbEpisode]], - ): + ) -> None: """ - 添加原始对象引用 + 以双下划线键名将原始对象引用写入 ``context``。 + + 约定：消费方仅读不写，避免在事件回调里改这些对象污染下游流程。 """ raw_objects = { # 文件元数据 @@ -282,7 +313,7 @@ class TemplateContextBuilder: # 当前季的全部集信息 "__episodes_info__": episodes_info, } - self._context.update(raw_objects) + context.update(raw_objects) @staticmethod def __convert_invalid_characters(filename: str): diff --git a/app/modules/filemanager/transhandler.py b/app/modules/filemanager/transhandler.py index 277f6a0c..1bea26a1 100644 --- a/app/modules/filemanager/transhandler.py +++ b/app/modules/filemanager/transhandler.py @@ -20,6 +20,7 @@ from app.schemas import ( FileItem, TransferInterceptEventData, TransferOverwriteCheckEventData, + TransferRenameBuildEventData, TransferRenameEventData, ) from app.schemas.types import MediaType, ChainEventType @@ -1142,6 +1143,19 @@ class TransHandler: :param source_item: 源文件信息，即待整理的文件信息 :return: 生成的完整路径 """ + # 渲染前先发事件，让插件有机会往 rename_dict 写字段 + build_event_data = TransferRenameBuildEventData( + template_string=template_string, + rename_dict=rename_dict, + source_path=source_path, + source_item=source_item, + ) + build_event = eventmanager.send_event( + ChainEventType.TransferRenameBuild, build_event_data + ) + if build_event and build_event.event_data: + rename_dict = build_event.event_data.rename_dict + # 创建jinja2模板对象 template = Template(template_string) # 渲染生成的字符串 diff --git a/app/schemas/event.py b/app/schemas/event.py index fb0592ac..6d20f8f1 100644 --- a/app/schemas/event.py +++ b/app/schemas/event.py @@ -183,6 +183,44 @@ class CommandRegisterEventData(ChainEventData): source: str = Field(default="未知拦截源", description="拦截源") +class TransferRenameBuildEventData(ChainEventData): + """ + TransferRenameBuild 事件的数据模型 + + 在 ``transhandler.get_rename_path`` 渲染文件名之前发出，给插件一次往 + ``rename_dict`` 写字段的机会。典型用法是通过 ffprobe 或外部接口探测源文件， + 把分辨率、视频/音频编码、HDR 等字段写入 ``rename_dict``，主程序下一步渲染时 + 就能直接用到这些字段，不需要插件事后再渲染一次去覆盖结果。 + + 与 ``TransferRenameEventData`` 的分工： + - 本事件负责"往 ``rename_dict`` 里写字段"，没有输出参数； + - ``TransferRename`` 在渲染之后触发，负责对已渲染好的字符串再做改写（大小写、 + 词替换、模板覆盖等），由智能重命名一类插件使用。 + + 使用约定： + - 只往 ``rename_dict`` 写字段，不要在这里改写已经渲染好的字符串； + - ``source_path`` / ``source_item`` 为空时（如重命名预览场景），需要源文件 + 才能工作的插件请直接 return； + - ``rename_dict`` 中以双下划线开头的键（``__meta__`` / ``__mediainfo__`` 等） + 存放的是原始对象引用，只读使用，不要修改这些对象本身。 + + Attributes: + template_string (str): Jinja2 模板字符串 + rename_dict (Dict[str, Any]): 渲染上下文，可直接修改 + source_path (Optional[str]): 源文件路径，即待整理的文件路径 + source_item (Optional[FileItem]): 源文件信息，即待整理的文件信息 + """ + + template_string: str = Field(..., description="模板字符串") + rename_dict: Dict[str, Any] = Field(..., description="渲染上下文") + source_path: Optional[str] = Field( + None, description="源文件路径，即待整理的文件路径" + ) + source_item: Optional[FileItem] = Field( + None, description="源文件信息，即待整理的文件信息" + ) + + class TransferRenameEventData(ChainEventData): """ TransferRename 事件的数据模型 diff --git a/app/schemas/types.py b/app/schemas/types.py index 38ade1c6..e249503b 100644 --- a/app/schemas/types.py +++ b/app/schemas/types.py @@ -154,6 +154,8 @@ class ChainEventType(Enum): CommandRegister = "command.register" # 整理重命名 TransferRename = "transfer.rename" + # 整理重命名上下文构建 + TransferRenameBuild = "transfer.rename.build" # 整理拦截 TransferIntercept = "transfer.intercept" # 整理覆盖检查 diff --git a/tests/test_template_context_builder.py b/tests/test_template_context_builder.py new file mode 100644 index 00000000..b0c9442d --- /dev/null +++ b/tests/test_template_context_builder.py @@ -0,0 +1,72 @@ +#!/usr/bin/env python +# -*- coding:utf-8 -*- +""" +TemplateContextBuilder 的并发安全单元测试。 + +历史上 builder 持有 ``self._context`` 实例字段，``build()`` 内 ``clear()`` → +``_add_*`` → 推导式返回这一序列在 ``TRANSFER_THREADS > 1`` 下会被多线程相互 +覆盖，导致同一 builder 实例并发调用产生互相串味的 rename_dict。本测试在多 +线程下连续调用 ``build()``，校验每个线程拿到的字典只反映自己的入参。 +""" +import threading +import unittest + +from app.helper.message import TemplateContextBuilder + + +class TemplateContextBuilderConcurrencyTest(unittest.TestCase): + """ + 使用 8 个线程并发调用同一 TemplateContextBuilder 实例的 build()， + 确保各自的 file_extension / 自定义 kwargs 不会被其它线程覆盖。 + """ + + THREAD_COUNT = 8 + ITERATIONS_PER_THREAD = 200 + + def test_concurrent_build_no_cross_contamination(self): + builder = TemplateContextBuilder() + errors = [] + + def worker(tag: int) -> None: + try: + for _ in range(self.ITERATIONS_PER_THREAD): + ctx = builder.build( + file_extension=f".{tag}", + marker=tag, + ) + self.assertEqual(ctx.get("fileExt"), f".{tag}") + self.assertEqual(ctx.get("marker"), tag) + except AssertionError as exc: + errors.append(exc) + + threads = [ + threading.Thread(target=worker, args=(i,), name=f"builder-{i}") + for i in range(self.THREAD_COUNT) + ] + for t in threads: + t.start() + for t in threads: + t.join() + + self.assertFalse( + errors, + msg=f"检测到并发串味，共 {len(errors)} 条；首个错误：{errors[0] if errors else ''}", + ) + + def test_build_returns_independent_dicts(self): + """ + 即便不开线程，连续两次 build() 也应当返回相互独立的 dict 实例， + 避免无状态化后调用方误以为返回的还是 builder 内部共享对象。 + """ + builder = TemplateContextBuilder() + first = builder.build(file_extension=".a", marker=1) + second = builder.build(file_extension=".b", marker=2) + self.assertIsNot(first, second) + self.assertEqual(first.get("fileExt"), ".a") + self.assertEqual(second.get("fileExt"), ".b") + # 第二次调用不应反向污染第一次的结果 + self.assertEqual(first.get("marker"), 1) + + +if __name__ == "__main__": + unittest.main() diff --git a/tests/test_transfer_rename_build_event.py b/tests/test_transfer_rename_build_event.py new file mode 100644 index 00000000..12cc67e7 --- /dev/null +++ b/tests/test_transfer_rename_build_event.py @@ -0,0 +1,141 @@ +#!/usr/bin/env python +# -*- coding:utf-8 -*- +""" +TransferRenameBuild 事件的单元测试。 + +通过 patch ``eventmanager.send_event`` 模拟链上插件处理器，避免依赖 +MoviePilot 的"插件实例反查"机制（``__get_class_instance`` 要求 handler 位于 +``app.`` 模块内）。覆盖以下关键路径： + +1. 插件就地 mutate ``rename_dict`` 后，主程序首次渲染读取到补充字段； +2. 插件用"返回新 dict"的方式（替换 ``event_data.rename_dict`` 引用）也能生效； +3. 没有任何监听者时（``send_event`` 返回 None），渲染输出与改造前完全一致； +4. ``get_rename_path`` 把正确的 ``source_path`` / ``source_item`` / ``template_string`` + 注入到 ``TransferRenameBuildEventData``，便于插件做条件早退。 +""" +import unittest +from unittest.mock import patch + +from app.core.event import Event +from app.modules.filemanager.transhandler import TransHandler +from app.schemas.event import TransferRenameBuildEventData +from app.schemas.types import ChainEventType + + +class TransferRenameBuildEventTest(unittest.TestCase): + """ + 通过 ``unittest.mock.patch`` 替换 ``eventmanager.send_event`` 实现： + 每个测试自定义一个 ``fake_send_event``，根据事件类型决定如何模拟链上插件 + 对 ``event_data`` 的修改，再返回 ``Event(event_data=...)``，与真实 dispatcher + 返回行为对齐。 + """ + + TEMPLATE = "{{title}}{% if effect %} {{effect}}{% endif %}{% if codec %} {{codec}}{% endif %}" + + @staticmethod + def _make_event(event_type: ChainEventType, data) -> Event: + """构造真实 dispatcher 返回的 Event 对象，便于和生产代码读取路径对齐。""" + return Event(event_type=event_type.value, event_data=data) + + def test_in_place_field_supplement_takes_effect(self): + captured = {} + + def fake_send_event(event_type, data, **_kwargs): + if event_type is ChainEventType.TransferRenameBuild: + self.assertIsInstance(data, TransferRenameBuildEventData) + captured["template_string"] = data.template_string + captured["source_path"] = data.source_path + # 模拟插件就地补充字段 + data.rename_dict["effect"] = "SDR" + return self._make_event(event_type, data) + return self._make_event(event_type, data) + + with patch( + "app.modules.filemanager.transhandler.eventmanager.send_event", + side_effect=fake_send_event, + ): + path = TransHandler.get_rename_path( + template_string=self.TEMPLATE, + rename_dict={"title": "Foo"}, + source_path="/downloads/foo.mkv", + ) + + self.assertEqual(path.as_posix(), "Foo SDR") + self.assertEqual(captured["template_string"], self.TEMPLATE) + self.assertEqual(captured["source_path"], "/downloads/foo.mkv") + + def test_returning_new_dict_reference_is_respected(self): + """ + 模拟插件用"完整替换 rename_dict 引用"的写法，验证 get_rename_path 在事件 + 返回后会重新取引用，新 dict 中的字段也能被首次渲染读到。 + """ + + def fake_send_event(event_type, data, **_kwargs): + if event_type is ChainEventType.TransferRenameBuild: + new_dict = dict(data.rename_dict) + new_dict["codec"] = "H265" + data.rename_dict = new_dict + return self._make_event(event_type, data) + + with patch( + "app.modules.filemanager.transhandler.eventmanager.send_event", + side_effect=fake_send_event, + ): + path = TransHandler.get_rename_path( + template_string=self.TEMPLATE, + rename_dict={"title": "Foo"}, + source_path="/downloads/foo.mkv", + ) + + self.assertEqual(path.as_posix(), "Foo H265") + + def test_no_listeners_yields_unchanged_render(self): + """ + 监听者缺席时 send_event 返回 None；get_rename_path 应跳过引用刷新并按原 + rename_dict 渲染，行为与改造前完全一致。 + """ + + def fake_send_event(event_type, _data, **_kwargs): + # 真实 dispatcher 在无 enabled handler 时返回 None + return None + + with patch( + "app.modules.filemanager.transhandler.eventmanager.send_event", + side_effect=fake_send_event, + ): + path = TransHandler.get_rename_path( + template_string=self.TEMPLATE, + rename_dict={"title": "Foo"}, + source_path="/downloads/foo.mkv", + ) + + self.assertEqual(path.as_posix(), "Foo") + + def test_event_data_carries_source_metadata(self): + """ + 即便没有 source_path（recommend_name 预览场景），事件仍会触发，但 + ``source_path`` / ``source_item`` 都为 None，供插件自行早退。 + """ + captured = {} + + def fake_send_event(event_type, data, **_kwargs): + if event_type is ChainEventType.TransferRenameBuild: + captured["source_path"] = data.source_path + captured["source_item"] = data.source_item + return self._make_event(event_type, data) + + with patch( + "app.modules.filemanager.transhandler.eventmanager.send_event", + side_effect=fake_send_event, + ): + TransHandler.get_rename_path( + template_string=self.TEMPLATE, + rename_dict={"title": "Foo"}, + ) + + self.assertIsNone(captured["source_path"]) + self.assertIsNone(captured["source_item"]) + + +if __name__ == "__main__": + unittest.main()