"""一次性脚本:将外置模块中的 `def f(self, ...)` 自由函数转换为 Mixin 方式。 操作: 1. 给所有顶层 `def f(self, ...)` 加 `self: "PQAutomationApp"` 注解(仅注解,不移动)。 2. 在文件顶部(首个 def/class 之前、import 块之后)插入 TYPE_CHECKING 块。 3. 在文件末尾追加 `class XxxMixin:` 把这些函数作为类属性挂上。 不会改变: - 函数体(包括内部 `_xxx(self, ...)` 直接调用)。 - 已存在的类、模块级常量。 """ from __future__ import annotations import ast import io import os import re import sys import tokenize from dataclasses import dataclass # 文件 -> Mixin 类名 TARGETS: dict[str, str] = { "app/config_io.py": "ConfigIOMixin", "app/views/chart_frame.py": "ChartFrameMixin", "app/views/panels/main_layout.py": "MainLayoutMixin", "app/views/panels/cct_panel.py": "CctPanelMixin", "app/device/connection.py": "DeviceConnectionMixin", "app/views/panels/custom_template_panel.py": "CustomTemplatePanelMixin", "app/views/panels/side_panels.py": "SidePanelsMixin", "app/views/panels/ai_image_panel.py": "AIImagePanelMixin", "app/views/panels/single_step_panel.py": "SingleStepPanelMixin", "app/views/panels/pantone_baseline_panel.py": "PantoneBaselinePanelMixin", "app/tests/local_dimming.py": "LocalDimmingMixin", "app/views/panel_manager.py": "PanelManagerMixin", "app/runner/test_runner.py": "TestRunnerMixin", "app/plots/plot_gamut.py": "PlotGamutMixin", "app/plots/plot_gamma.py": "PlotGammaMixin", "app/plots/plot_eotf.py": "PlotEotfMixin", "app/plots/plot_cct.py": "PlotCctMixin", "app/plots/plot_contrast.py": "PlotContrastMixin", "app/plots/plot_accuracy.py": "PlotAccuracyMixin", } TYPE_CHECKING_BLOCK = ( "from typing import TYPE_CHECKING\n" "\n" "if TYPE_CHECKING:\n" " from pqAutomationApp import PQAutomationApp\n" "\n" ) @dataclass class SelfFunc: name: str lineno: int # 1-based, line of `def` col_offset: int end_lineno: int def_line_idx: int # 0-based line index of `def ...` line self_token_line_idx: int # 0-based line index of `self` self_token_col: int def _read(path: str) -> str: with open(path, "rb") as f: data = f.read() # 去 BOM if data.startswith(b"\xef\xbb\xbf"): data = data[3:] return data.decode("utf-8") def _write(path: str, text: str) -> None: with open(path, "w", encoding="utf-8", newline="\n") as f: f.write(text) def _locate_self_token(src_lines: list[str], def_line_idx: int) -> tuple[int, int] | None: """在 def 行(可能多行签名)中定位首个参数 `self` 的位置。""" # 用 tokenize 解析从 def 行开始的片段 snippet = "".join(src_lines[def_line_idx:]) try: tokens = list(tokenize.generate_tokens(io.StringIO(snippet).readline)) except tokenize.TokenizeError: return None saw_open_paren = False for tok in tokens: if tok.type == tokenize.OP and tok.string == "(": saw_open_paren = True continue if saw_open_paren and tok.type == tokenize.NAME and tok.string == "self": # tok.start 是 (row, col) 相对于 snippet(1-based row) row_in_snippet = tok.start[0] - 1 col = tok.start[1] return (def_line_idx + row_in_snippet, col) if saw_open_paren and tok.type == tokenize.OP and tok.string in (")", ","): # 第一个参数不是 self —— 跳过 return None return None def collect_self_funcs(src: str) -> tuple[ast.Module, list[SelfFunc]]: tree = ast.parse(src) src_lines = src.splitlines(keepends=True) results: list[SelfFunc] = [] for node in tree.body: if isinstance(node, ast.FunctionDef) and node.args.args and node.args.args[0].arg == "self": def_idx = node.lineno - 1 pos = _locate_self_token(src_lines, def_idx) if pos is None: continue results.append(SelfFunc( name=node.name, lineno=node.lineno, col_offset=node.col_offset, end_lineno=node.end_lineno or node.lineno, def_line_idx=def_idx, self_token_line_idx=pos[0], self_token_col=pos[1], )) return tree, results def annotate_self(src: str, funcs: list[SelfFunc]) -> str: """把每个 def 的首个 `self` 形参替换为 `self: "PQAutomationApp"`。""" lines = src.splitlines(keepends=True) # 从后往前替换,避免行号变动 for fn in sorted(funcs, key=lambda f: -f.self_token_line_idx): line = lines[fn.self_token_line_idx] col = fn.self_token_col # 检查后续是否已经有注解 after = line[col + len("self"):] # 已经注解过则跳过 m = re.match(r"\s*:", after) if m: continue new_line = line[:col] + 'self: "PQAutomationApp"' + line[col + len("self"):] lines[fn.self_token_line_idx] = new_line return "".join(lines) def ensure_type_checking_block(src: str) -> str: if "from pqAutomationApp import PQAutomationApp" in src: return src # 找到首个非 docstring / 非注释 / 非 import 的位置: # 简单策略:在最后一个 import 行之后插入;若没有 import,则在 docstring 之后插入。 lines = src.splitlines(keepends=True) insert_idx = 0 in_docstring = False doc_quote: str | None = None last_import_idx = -1 for i, line in enumerate(lines): stripped = line.lstrip() if i == 0 and (stripped.startswith('"""') or stripped.startswith("'''")): q = stripped[:3] doc_quote = q if stripped.count(q) >= 2 and len(stripped) > 3: # 单行 docstring last_import_idx = max(last_import_idx, i) continue in_docstring = True continue if in_docstring: if doc_quote and doc_quote in line: in_docstring = False last_import_idx = max(last_import_idx, i) continue if stripped.startswith("import ") or stripped.startswith("from "): last_import_idx = i continue if stripped == "" or stripped.startswith("#"): continue # 遇到第一个真实代码 break insert_idx = last_import_idx + 1 new_lines = lines[:insert_idx] + ["\n", TYPE_CHECKING_BLOCK] + lines[insert_idx:] return "".join(new_lines) def append_mixin(src: str, mixin_name: str, func_names: list[str]) -> str: if f"class {mixin_name}" in src: return src body_lines = [] body_lines.append("") body_lines.append("") body_lines.append(f"class {mixin_name}:") body_lines.append(f' """由 tools/refactor_to_mixins.py 自动生成。') body_lines.append(" 把本模块的自由函数挂到 PQAutomationApp 上,便于 F12 跳转与类型推断。") body_lines.append(' """') for name in func_names: body_lines.append(f" {name} = {name}") text = "\n".join(body_lines) + "\n" if not src.endswith("\n"): src += "\n" return src + text def process(path: str, mixin_name: str) -> None: src = _read(path) tree, funcs = collect_self_funcs(src) if not funcs: print(f" -> skip (no self-funcs)") return func_names = [f.name for f in funcs] new_src = annotate_self(src, funcs) new_src = ensure_type_checking_block(new_src) new_src = append_mixin(new_src, mixin_name, func_names) _write(path, new_src) print(f" -> {mixin_name} with {len(func_names)} methods") def main(): root = os.path.dirname(os.path.dirname(os.path.abspath(__file__))) os.chdir(root) for rel, mixin in TARGETS.items(): print(f"Processing {rel}") process(rel, mixin) if __name__ == "__main__": main()