pqAutomationApp/tools/refactor_to_mixins.py

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