app.py

"""
interactive_modify_figures_from_python - ggplot 图形交互式编辑器
加载目录下保存的 .rds ggplot 对象,网页端实时修改文字/样式/几何参数,
改动以补丁(patches.json)形式保存,可随时重出 pdf/svg/png/tiff。
后端通过常驻 R 进程(localhost socket)渲染,避免每次重启 R 的开销。
"""

import hashlib
import json
import os
import re
import socket
import subprocess
import threading
import time

from flask import Flask, jsonify, render_template, request, send_from_directory

BASE_DIR = os.path.dirname(os.path.abspath(__file__))
RENDER_DIR = os.path.join(BASE_DIR, "renders")
FIG_ROOT = os.environ.get(
    "FIG_ROOT", "/path/to/your/ggplot_figures"
)
RSCRIPT = os.environ.get("RSCRIPT", "/usr/bin/Rscript")
HTTP_PORT = int(os.environ.get("FIG_EDITOR_PORT", "5081"))

os.makedirs(RENDER_DIR, exist_ok=True)

app = Flask(__name__)


# ── R 工作进程管理 ──────────────────────────────────────────────────────────
class RWorker:
    """常驻 R 进程,localhost socket,**每请求一个新连接**(R 端回完即关)。
    不能用 stdin/stdout 管道:readRDS 会干扰后续 stdin 读取。
    长连接曾出现两端状态错位(请求发进死连接,挂到超时),故每请求新建连接。"""

    def __init__(self):
        self.proc = None
        self.port = None
        self.lock = threading.Lock()
        self.start()

    def start(self):
        self.stop()
        # 清掉残留 worker(如 Flask 被 kill 后遗留的孤儿进程)。
        # 字符类 [.] 避免 pgrep 模式匹配到这条 bash 命令自身。
        subprocess.run(
            'pgrep -f "r_worker[.]R" | xargs -r kill', shell=True, check=False
        )
        time.sleep(0.5)
        # 每次用随机空闲端口:固定端口在上个 worker 被杀后会进入 TIME_WAIT,
        # R 的 socketConnection 没有 SO_REUSEADDR,短时间内 bind 必失败。
        probe = socket.socket()
        probe.bind(("127.0.0.1", 0))
        self.port = probe.getsockname()[1]
        probe.close()
        env = dict(os.environ, FIG_WORKER_PORT=str(self.port))
        log = open(os.path.join(BASE_DIR, "r_worker.log"), "ab")
        self.proc = subprocess.Popen(
            [RSCRIPT, os.path.join(BASE_DIR, "r_worker.R")],
            stdout=subprocess.DEVNULL, stderr=log, env=env,
        )
        # 等 worker 监听就绪(首次加载包要几秒)
        deadline = time.time() + 60
        while time.time() < deadline:
            if self.proc.poll() is not None:
                raise RuntimeError("R worker 启动失败,见 r_worker.log")
            try:
                s = socket.create_connection(("127.0.0.1", self.port), timeout=2)
                s.close()  # 占位 accept:R 端读 0 行后关闭,无副作用
                print(f"[worker] ready on port {self.port}", flush=True)
                return
            except OSError:
                time.sleep(0.3)
        raise RuntimeError("R worker 启动超时")

    def stop(self):
        if self.proc and self.proc.poll() is None:
            try:
                self.proc.terminate()
                self.proc.wait(timeout=5)
            except Exception:
                self.proc.kill()
        self.proc = None

    def request(self, payload, timeout=120, _retry=True):
        with self.lock:
            try:
                if self.proc is None or self.proc.poll() is not None:
                    print("[worker] 已退出,重启", flush=True)
                    self.start()
                s = socket.create_connection(("127.0.0.1", self.port), timeout=10)
                s.settimeout(timeout)
                f = s.makefile("rw")
                f.write(json.dumps(payload, ensure_ascii=False) + "\n")
                f.flush()
                resp = f.readline()
                try:
                    f.close()
                    s.close()
                except OSError:
                    pass
                if not resp:
                    raise ConnectionError("worker 空响应")
                return json.loads(resp)
            except (OSError, ConnectionError, json.JSONDecodeError) as e:
                print(f"[worker] 请求失败({type(e).__name__}: {e}),cmd={payload.get('cmd')}", flush=True)
                if not _retry:
                    return {"ok": False, "error": f"R 进程通讯失败: {e}"}
                try:
                    self.start()
                except Exception as e2:
                    return {"ok": False, "error": f"R 进程重启失败: {e2}"}
                return self.request(payload, timeout, _retry=False)


worker = RWorker()


# ── 图形清单 ────────────────────────────────────────────────────────────────
_fig_cache = {"t": 0, "data": None}


def _group_key(rel_dir):
    """面板分组键:去掉路径里所有 vN 版本目录段。
    图2/v2/b_macro_auc/v3/v4 -> 图2/b_macro_auc"""
    return "/".join(s for s in rel_dir.split(os.sep) if not re.fullmatch(r"v\d+", s))


def _dir_version(rel_dir):
    """目录的最高版本号,无 vN 段为 0。"""
    vs = [int(m.group(1)) for s in rel_dir.split(os.sep)
          for m in [re.fullmatch(r"v(\d+)", s)] if m]
    return max(vs) if vs else 0


def list_figures():
    if time.time() - _fig_cache["t"] < 30 and _fig_cache["data"] is not None:
        return _fig_cache["data"]
    figs = []
    for dirpath, _, files in os.walk(FIG_ROOT):
        for f in sorted(files):
            if f.endswith(".rds"):
                full = os.path.join(dirpath, f)
                rel = os.path.relpath(full, FIG_ROOT)
                rel_dir = os.path.dirname(rel)
                figs.append({
                    "id": rel,
                    "label": rel[:-4],
                    "group": rel.split(os.sep)[0],
                    "mtime": os.path.getmtime(full),
                    "key": _group_key(rel_dir),
                    "version": _dir_version(rel_dir),
                })
    # 每组只把最高版本标为 latest(默认只显示这些)
    best = {}
    for f in figs:
        best[f["key"]] = max(best.get(f["key"], -1), f["version"])
    for f in figs:
        f["latest"] = f["version"] == best[f["key"]]
    figs.sort(key=lambda x: x["id"])
    _fig_cache.update(t=time.time(), data=figs)
    return figs


def fig_path(rel):
    full = os.path.realpath(os.path.join(FIG_ROOT, rel))
    if not full.startswith(os.path.realpath(FIG_ROOT)) or not os.path.isfile(full):
        raise FileNotFoundError(rel)
    return full


# ── 路由 ────────────────────────────────────────────────────────────────────
@app.route("/")
def index():
    return render_template("index.html")


@app.route("/api/figures")
def api_figures():
    return jsonify({"ok": True, "figures": list_figures(), "root": FIG_ROOT})


@app.route("/api/inspect")
def api_inspect():
    rel = request.args.get("fig", "")
    subplot = int(request.args.get("subplot", "1"))
    try:
        path = fig_path(rel)
    except FileNotFoundError:
        return jsonify({"ok": False, "error": "文件不存在"})
    res = worker.request({"cmd": "inspect", "path": path, "subplot": subplot})
    # 附带已保存的补丁
    patches_file = os.path.splitext(path)[0] + ".patches.json"
    if os.path.isfile(patches_file):
        try:
            res["saved"] = json.load(open(patches_file, encoding="utf-8"))
        except Exception:
            pass
    # 原始画布尺寸:从同目录 png 的像素/dpi 反推英寸(rds 里不存 ggsave 尺寸)
    png_file = os.path.splitext(path)[0] + ".png"
    if os.path.isfile(png_file):
        try:
            from PIL import Image
            im = Image.open(png_file)
            dpi = (im.info.get("dpi") or (300, 300))[0] or 300
            res["orig_canvas"] = {
                "width": round(im.size[0] / dpi, 2),
                "height": round(im.size[1] / dpi, 2),
            }
        except Exception:
            pass
    return jsonify(res)


def render_name(rel, subplot):
    h = hashlib.md5(f"{rel}#{subplot}".encode()).hexdigest()[:12]
    return f"{h}.svg"


@app.route("/api/render", methods=["POST"])
def api_render():
    body = request.get_json(force=True)
    rel = body.get("fig", "")
    subplot = int(body.get("subplot", 1))
    patches = body.get("patches", {})
    patches_all = body.get("patches_all") or {str(subplot): patches}
    moves = body.get("moves", {})
    width = float(body.get("width", 10))
    height = float(body.get("height", 6))
    try:
        path = fig_path(rel)
    except FileNotFoundError:
        return jsonify({"ok": False, "error": "文件不存在"})
    name = render_name(rel, 0)  # 整图渲染,文件名与子图无关
    out = os.path.join(RENDER_DIR, name)
    res = worker.request({
        "cmd": "render", "path": path, "subplot": subplot,
        "patches": patches, "patches_all": patches_all, "moves": moves,
        "width": width, "height": height, "out": out,
    })
    if res.get("ok"):
        # 相对路径:经反向代理(子路径前缀)访问时绝对路径会丢前缀
        res["url"] = f"renders/{name}?v={time.time()}"
        res["code"] = gen_r_code(patches_all, moves)
    return jsonify(res)


@app.route("/api/save", methods=["POST"])
def api_save():
    body = request.get_json(force=True)
    rel = body.get("fig", "")
    patches_all = body.get("patches_all", {})
    width = float(body.get("width", 10))
    height = float(body.get("height", 6))
    try:
        path = fig_path(rel)
    except FileNotFoundError:
        return jsonify({"ok": False, "error": "文件不存在"})
    outdir = os.path.dirname(path)
    prefix = os.path.splitext(os.path.basename(path))[0]
    moves = body.get("moves", {})
    res = worker.request({
        "cmd": "save", "path": path, "patches_all": patches_all, "moves": moves,
        "width": width, "height": height, "outdir": outdir, "prefix": prefix,
    })
    if res.get("ok"):
        # 等效 R 代码也落盘,方便粘回原脚本
        code_lines = [f'p <- readRDS("{path}")']
        for sp, patches in sorted(patches_all.items(), key=lambda kv: int(kv[0])):
            code_lines += gen_r_patch_lines(patches, int(sp))
        for k, mv in moves.items():
            code_lines.append(f'# 移动文本 {k}: dx={mv.get("dx")}pt, dy={mv.get("dy")}pt(见 patches.json)')
        code_lines.append(f'# 然后 ggsave("{prefix}.pdf", p, width={width}, height={height}) 等')
        code_file = os.path.join(outdir, prefix + ".patches.R")
        with open(code_file, "w", encoding="utf-8") as f:
            f.write("\n".join(code_lines) + "\n")
        res["files"] = res.get("files", []) + [os.path.basename(code_file)]
    return jsonify(res)


@app.route("/renders/<path:name>")
def renders(name):
    resp = send_from_directory(RENDER_DIR, name)
    resp.headers["Cache-Control"] = "no-store"
    return resp


# ── 等效 R 代码生成(展示用) ────────────────────────────────────────────────
def r_quote(s):
    return '"' + str(s).replace("\\", "\\\\").replace('"', '\\"') + '"'


def gen_r_patch_lines(patches, subplot=None):
    lines = []
    yl = {}
    for pid, val in patches.items():
        parts = pid.split("|")
        kind = parts[0]
        if kind == "label":
            lines.append(f"p$labels${parts[1]} <- {r_quote(val)}")
        elif kind == "text":
            el, prop = parts[1], parts[2]
            v = val if prop in ("size", "hjust", "vjust", "angle") else r_quote(val)
            lines.append(f"p <- p + theme({el} = element_text({prop} = {v}))")
        elif kind == "theme" and parts[1] == "legend.position":
            lines.append(f"p <- p + theme(legend.position = {r_quote(val)})")
        elif kind == "global" and parts[1] == "text.size":
            lines.append(f"p <- p + theme(text = element_text(size = {val}))")
        elif kind == "layer":
            i, slot, param = parts[1], parts[2], "|".join(parts[3:])
            v = r_quote(val) if isinstance(val, str) else json.dumps(val)
            lines.append(f"p$layers[[{i}]]${slot}${param} <- {v}")
        elif kind == "palette":
            i, key = parts[1], "|".join(parts[2:])
            lines.append(f'# 配色: scale {i} 的 {r_quote(key)} <- {r_quote(val)}(见 patches.json)')
        elif kind == "scalelabel":
            i, key = parts[1], "|".join(parts[2:])
            lines.append(f'# 刻度文字: scale {i} 的 {r_quote(key)} <- {r_quote(val)}(见 patches.json)')
        elif kind == "strip":
            lines.append(f'# 分面标签: {r_quote(parts[1])} <- {r_quote(val)}(见 patches.json)')
        elif kind == "ylim":
            yl[parts[1]] = val
    if yl:
        lines.append(f"p$coordinates$params$limits$y <- c({yl.get('min')}, {yl.get('max')})")
    if subplot is not None and subplot > 1:
        lines.insert(0, f"# —— 子图 {subplot} ——")
    return lines


def gen_r_code(patches_all, moves):
    lines = []
    for sp, patches in sorted(patches_all.items(), key=lambda kv: int(kv[0])):
        lines += gen_r_patch_lines(patches, int(sp) if len(patches_all) > 1 else None)
    for k, mv in moves.items():
        lines.append(f'# 移动文本 {k}: dx={mv.get("dx")}pt, dy={mv.get("dy")}pt(见 patches.json)')
    return "\n".join(lines) if lines else "# 尚未修改"


def warmup():
    """后台预热:启动后把最新版本图的 inspect 全部跑一遍(结果进 worker 缓存),
    用户切换图时不再等待 readRDS + ggplot_build。"""
    time.sleep(2)
    try:
        for f in list_figures():
            if not f.get("latest"):
                continue
            try:
                path = fig_path(f["id"])
            except FileNotFoundError:
                continue
            res = worker.request({"cmd": "inspect", "path": path, "subplot": 1})
            # patchwork 的其余子图也预热
            for sp in (res.get("subplots") or [])[1:]:
                worker.request({"cmd": "inspect", "path": path,
                                "subplot": sp["index"]})
    except Exception:
        pass


if __name__ == "__main__":
    threading.Thread(target=warmup, daemon=True).start()
    app.run(host="0.0.0.0", port=HTTP_PORT, threaded=True)