225 lines
7.9 KiB
Python
225 lines
7.9 KiB
Python
"""一次性脚本：将外置模块中的 `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()
|