autotrain/scripts/record_step.py

#!/usr/bin/env python3
"""录制一次"用户操作"前后的手机状态。

用法:
  python3 scripts/record_step.py "label"            # 记一次状态,自动序号
  python3 scripts/record_step.py --new "first label" # 开新 session,清重置
  python3 scripts/record_step.py --list             # 列出当前 session 已记的步骤

每次调用会在 work/replay/<session_ts>/<NN_label>/ 下产出:
  - ui.xml          adb uiautomator dump
  - screen.png      adb 截屏
  - meta.json       时间戳 + label + 屏幕分辨率
  - texts.txt       从 xml 抽出的所有可见 text(后续 diff 用)
"""

from __future__ import annotations

import argparse
import json
import re
import subprocess
import sys
import time
from pathlib import Path

import device_runtime

PROJECT_ROOT = Path(__file__).resolve().parents[1]
REPLAY_ROOT = PROJECT_ROOT / "work" / "replay"
CURRENT_SESSION_POINTER = REPLAY_ROOT / "_current_session"


def latest_session() -> Path | None:
    if CURRENT_SESSION_POINTER.exists():
        target = Path(CURRENT_SESSION_POINTER.read_text().strip())
        if target.exists():
            return target
    return None


def new_session() -> Path:
    REPLAY_ROOT.mkdir(parents=True, exist_ok=True)
    ts = time.strftime("%Y%m%d_%H%M%S")
    session = REPLAY_ROOT / ts
    session.mkdir(parents=True, exist_ok=False)
    CURRENT_SESSION_POINTER.write_text(str(session))
    return session


def next_step_index(session: Path) -> int:
    existing = sorted(p.name for p in session.iterdir() if p.is_dir())
    n = 0
    for name in existing:
        m = re.match(r"(\d+)_", name)
        if m:
            n = max(n, int(m.group(1)))
    return n + 1


def slugify(label: str) -> str:
    label = re.sub(r"\s+", "_", label.strip())
    return re.sub(r"[^a-zA-Z0-9_一-鿿-]+", "", label)[:60] or "step"


def dump_ui_xml(target: Path) -> None:
    device_runtime.run_adb(["shell", "uiautomator", "dump", "/sdcard/_record.xml"])
    device_runtime.run_adb(["pull", "/sdcard/_record.xml", str(target)])


def take_screenshot(target: Path) -> None:
    result = device_runtime.run_adb_bytes(["exec-out", "screencap", "-p"])
    target.write_bytes(result.stdout)


def extract_texts(xml_path: Path) -> list[str]:
    xml = xml_path.read_text(encoding="utf-8", errors="ignore")
    return [t for t in re.findall(r'text="([^"]+)"', xml) if t.strip()]


def png_dimensions(path: Path) -> tuple[int, int]:
    data = path.read_bytes()
    if data[:8] == b"\x89PNG\r\n\x1a\n":
        return int.from_bytes(data[16:20], "big"), int.from_bytes(data[20:24], "big")
    return 0, 0


def record(session: Path, label: str) -> Path:
    idx = next_step_index(session)
    folder = session / f"{idx:02d}_{slugify(label)}"
    folder.mkdir(parents=True, exist_ok=False)

    xml_path = folder / "ui.xml"
    png_path = folder / "screen.png"
    dump_ui_xml(xml_path)
    take_screenshot(png_path)

    texts = extract_texts(xml_path)
    (folder / "texts.txt").write_text("\n".join(texts), encoding="utf-8")

    w, h = png_dimensions(png_path)
    meta = {
        "label": label,
        "index": idx,
        "timestamp": time.strftime("%Y-%m-%d %H:%M:%S"),
        "screen_w": w,
        "screen_h": h,
        "text_count": len(texts),
    }
    (folder / "meta.json").write_text(json.dumps(meta, ensure_ascii=False, indent=2), encoding="utf-8")
    return folder


def cmd_list(session: Path) -> None:
    for folder in sorted(session.iterdir()):
        if folder.is_dir():
            meta_path = folder / "meta.json"
            if meta_path.exists():
                meta = json.loads(meta_path.read_text())
                print(f"  {folder.name:30s} {meta.get('timestamp','')}  text_count={meta.get('text_count')}")


def main(argv: list[str] | None = None) -> int:
    parser = argparse.ArgumentParser(description="录制手机操作前后状态")
    parser.add_argument("label", nargs="?", help="本步骤的描述,作为目录名一部分")
    parser.add_argument("--new", dest="new", action="store_true", help="开新 session")
    parser.add_argument("--list", dest="list_only", action="store_true", help="列出当前 session 已记的步骤")
    args = parser.parse_args(argv)

    if args.list_only:
        session = latest_session()
        if not session:
            print("no_session")
            return 1
        print(f"session={session}")
        cmd_list(session)
        return 0

    if args.new or latest_session() is None:
        session = new_session()
        print(f"new_session={session}")
    else:
        session = latest_session()

    if not args.label:
        print("error: label required", file=sys.stderr)
        return 2

    folder = record(session, args.label)
    print(f"recorded={folder.name} session={session.name}")
    return 0


if __name__ == "__main__":
    raise SystemExit(main())