mirror of
https://github.com/jxxghp/MoviePilot.git
synced 2026-10-29 03:18:14 +08:00
78 lines
3.3 KiB
Python
78 lines
3.3 KiB
Python
"""
|
||
临时替换模块缓存并按快照恢复的测试工具。
|
||
|
||
stub_modules 在上下文中替换指定模块，snapshot_modules 和 restore_modules 可成对用于
|
||
测试夹具的准备与清理。所有操作只调整 sys.modules 映射，不会撤销导入期间的其他副作用，
|
||
也不会更新调用方已经绑定的对象引用或父包上的模块属性。
|
||
"""
|
||
import sys
|
||
from contextlib import contextmanager
|
||
from typing import Any, Callable, Dict, Iterator, Optional
|
||
|
||
|
||
@contextmanager
|
||
def stub_modules(stubs: Dict[str, Any]) -> Iterator[None]:
|
||
"""
|
||
在上下文中替换指定模块，退出时恢复原有模块对象。
|
||
|
||
原条目缺失或值为 None 时，退出后移除该键；其他原值按对象恢复。
|
||
仅恢复显式列出的模块，目标模块导入时产生的其他状态由调用方负责清理。
|
||
|
||
:param stubs: 模块全名到占位对象的映射，通常使用 ModuleType 或 MagicMock
|
||
:yields: 已替换模块的执行区间，不提供额外值
|
||
"""
|
||
saved: Dict[str, Any] = {}
|
||
for name, module in stubs.items():
|
||
saved[name] = sys.modules.get(name)
|
||
sys.modules[name] = module
|
||
try:
|
||
yield
|
||
finally:
|
||
for name, original in saved.items():
|
||
if original is None:
|
||
sys.modules.pop(name, None)
|
||
else:
|
||
sys.modules[name] = original
|
||
|
||
|
||
def snapshot_modules(prefix: Optional[str] = None) -> Dict[str, Any]:
|
||
"""
|
||
对指定范围内的模块缓存创建浅快照。
|
||
|
||
快照不复制模块内部状态，也不冻结模块对象属性。
|
||
|
||
:param prefix: 包名前缀；None 表示全部模块，末尾点号忽略，只匹配父包及其子模块
|
||
:return: 模块全名到原对象的映射，可交给 restore_modules 恢复
|
||
"""
|
||
if prefix is None:
|
||
return dict(sys.modules)
|
||
# 归一去掉末尾点后按"精确父模块或其子模块路径"匹配，避免 prefix="app" 误配到 "apple"
|
||
parent = prefix.rstrip(".")
|
||
return {k: v for k, v in sys.modules.items() if k == parent or k.startswith(parent + ".")}
|
||
|
||
|
||
def restore_modules(snapshot: Dict[str, Any], prefix: Optional[str] = None) -> None:
|
||
"""
|
||
恢复快照范围内的模块映射，并移除该范围内后来新增的条目。
|
||
|
||
仅恢复对象引用，不撤销模块内部属性变化或导入的其他副作用。
|
||
|
||
:param snapshot: 此前 snapshot_modules 返回的浅快照
|
||
:param prefix: 与取快照时一致的包名前缀；None 表示恢复全部模块
|
||
"""
|
||
in_scope: Callable[[str], bool]
|
||
if prefix is None:
|
||
# 两个分支统一保存一个局部判定函数，保留赋值式 lambda 的既有写法。
|
||
in_scope = lambda name: True # noqa: E731
|
||
else:
|
||
head = prefix.rstrip(".")
|
||
# 与快照使用相同的包边界；局部 lambda 统一后续遍历的范围判定。
|
||
in_scope = lambda name: name == head or name.startswith(head + ".") # noqa: E731
|
||
# 移除范围内、快照中没有的新增项（通常是测试塞入的假桩）
|
||
for name in [n for n in sys.modules if in_scope(n) and n not in snapshot]:
|
||
sys.modules.pop(name, None)
|
||
# 恢复范围内被替换/删除的项
|
||
for name, module in snapshot.items():
|
||
if in_scope(name) and sys.modules.get(name) is not module:
|
||
sys.modules[name] = module
|