use idea cli to format java project, only format changed lines. support both git/jujutsu repo, can be executed directly or git commit hook.

README.md

If you don't use idea, it's hard to maintain the same code style with your colleagues. This script solves this problem, it uses idea cli and the code style XML of your project to format your code.

fmt.py

#!/usr/bin/env python3
import os
import platform
import re
import shutil
import subprocess
import sys
import tempfile
import time

# ================= 配置区域 =================
STYLE_XML_PATH = r"C:\programs\work\xxx.xml"

IDEA_BIN_PATHS = {
    "Windows": r"C:\software\IDEA\bin\format.bat",
    "Darwin": "/Applications/IntelliJ IDEA.app/Contents/bin/format.sh",
    "Linux": "/opt/intellij-idea/bin/format.sh",
}

ALLOWED_REPOS = [
    "my-test-repo",
]
# ===========================================


# ================= 核心工具函数 =================
def get_platform_idea_bin():
    system = platform.system()
    path = IDEA_BIN_PATHS.get(system)
    if not path:
        print(f"[Error] Unsupported OS: {system}")
        sys.exit(1)
    if not os.path.exists(path):
        print(f"[Error] IDEA formatter not found at: {path}")
        sys.exit(1)
    return path


def parse_modified_lines(diff_text: str) -> set:
    """
    解析标准的 Git/jj Unified Diff 文本。
    提取出当前工作副本（也就是 Diff 中的 '+' 侧）被修改的行号集合。
    """
    modified_lines = set()
    current_line = 0
    in_hunk = False

    for line in diff_text.splitlines():
        if line.startswith("@@"):
            # 匹配格式如: @@ -10,4 +10,5 @@
            match = re.match(r"^@@ -\d+(?:,\d+)? \+(\d+)(?:,\d+)? @@", line)
            if match:
                current_line = int(match.group(1))
                in_hunk = True
        elif in_hunk:
            if line.startswith("+++") or line.startswith("---"):
                continue
            if line.startswith("+"):
                modified_lines.add(current_line)
                current_line += 1
            elif line.startswith(" "):
                current_line += 1
            elif line.startswith("-"):
                # 所在位置被删除了内容，说明该行（对应新文件中的位置）有变更影响
                modified_lines.add(current_line)
            elif line.startswith("\\"):
                pass
            else:
                in_hunk = False
    return modified_lines


# ================= VCS (版本控制) 抽象与实现 =================
class VCS:
    name = "Unknown"

    def get_repo_root(self) -> str:
        raise NotImplementedError

    def get_modified_java_files(self) -> list:
        raise NotImplementedError

    def get_modified_lines(self, file_path: str) -> set:
        raise NotImplementedError

    def post_format_action(self, file_path: str) -> None:
        raise NotImplementedError


class GitVCS(VCS):
    name = "Git"

    def get_repo_root(self) -> str:
        return subprocess.check_output(
            ["git", "rev-parse", "--show-toplevel"], text=True
        ).strip()

    def get_modified_java_files(self) -> list:
        try:
            cmd = ["git", "diff", "--cached", "--name-only", "--diff-filter=ACMR"]
            result = subprocess.check_output(cmd, text=True, encoding="utf-8")
            return [
                f.strip() for f in result.splitlines() if f.strip().endswith(".java")
            ]
        except subprocess.CalledProcessError:
            print("[Error] Failed to get staged files via git.")
            sys.exit(1)

    def get_modified_lines(self, file_path: str) -> set:
        cmd = ["git", "diff", "--cached", "HEAD", "--", file_path]
        try:
            output = subprocess.check_output(
                cmd, encoding="utf-8", errors="ignore"
            ).strip()
            return parse_modified_lines(output)
        except Exception as e:
            print(f"  [Warning] Could not get diff lines for {file_path}: {e}")
            return set()

    def post_format_action(self, file_path: str) -> None:
        # Git 需要重新将格式化后的修改 stage 到暂存区
        subprocess.check_call(["git", "add", file_path])


class JjVCS(VCS):
    name = "Jujutsu (jj)"

    def get_repo_root(self) -> str:
        return subprocess.check_output(["jj", "workspace", "root"], text=True).strip()

    def get_modified_java_files(self) -> list:
        print(f"  -> [{self.name}] Fetching modified files summary...")
        try:
            # 打印带颜色的 summary 方便用户查看
            subprocess.run(["jj", "diff", "--summary", "--color=always"], check=False)
            print("-" * 40)

            cmd = ["jj", "diff", "--summary", "--color=never"]
            result = subprocess.check_output(
                cmd, text=True, encoding="utf-8", errors="ignore"
            )
            files = []
            for line in result.splitlines():
                line = line.strip()
                if not line:
                    continue
                parts = line.split(maxsplit=1)
                if len(parts) == 2:
                    status, filepath = parts
                    if status in ("A", "M", "C", "R") and filepath.endswith(".java"):
                        files.append(filepath)
            return files
        except subprocess.CalledProcessError as e:
            print(f"[Error] Failed to get modified files via jj. Error: {e}")
            sys.exit(1)

    def get_modified_lines(self, file_path: str) -> set:
        # jj 使用 diff --git 生成兼容 git 的 diff，方便复用解析器
        cmd = [
            "jj",
            "--no-pager",
            "diff",
            "--git",
            "-r",
            "@",
            file_path,
            "--color=never",
        ]
        try:
            output = subprocess.check_output(
                cmd, encoding="utf-8", errors="ignore"
            ).strip()
            return parse_modified_lines(output)
        except Exception as e:
            print(f"  [Warning] Could not get diff lines for {file_path}: {e}")
            return set()

    def post_format_action(self, file_path: str) -> None:
        pass  # jj 会自动追踪工作副本变更，不需要 add


def detect_vcs() -> VCS:
    """自动检测当前目录属于哪个版本控制系统 (优先检查 jj)"""
    if shutil.which("jj"):
        try:
            subprocess.check_output(["jj", "root"], stderr=subprocess.DEVNULL)
            return JjVCS()
        except subprocess.CalledProcessError:
            pass

    if shutil.which("git"):
        try:
            subprocess.check_output(
                ["git", "rev-parse", "--is-inside-work-tree"], stderr=subprocess.DEVNULL
            )
            return GitVCS()
        except subprocess.CalledProcessError:
            pass

    print("[Error] Neither Git nor jj repository detected, or CLI tools not installed.")
    sys.exit(1)


# ================= 格式化与 Patch 逻辑 =================
def run_idea_formatter_batch(idea_bin, target_dir):
    abs_style_path = os.path.abspath(STYLE_XML_PATH)
    cmd = [idea_bin, "-s", abs_style_path, "-r", target_dir]

    try:
        process = subprocess.Popen(
            cmd,
            stdout=subprocess.PIPE,
            stderr=subprocess.STDOUT,
            text=True,
            encoding="gbk" if platform.system() == "Windows" else "utf-8",
            errors="ignore",
        )

        start_time = time.time()
        while True:
            if time.time() - start_time > 120:
                print("  [Timeout] IDEA process took too long. Killing.")
                process.kill()
                break

            line = process.stdout.readline()
            if not line and process.poll() is not None:
                break
            if line:
                content = line.strip()
                if "ERROR" in content or "Formatted" in content:
                    print(f"[IDEA] {content}")

                if "file(s) formatted" in content or "files formatted" in content:
                    time.sleep(1)
                    process.terminate()
                    break

                if "Only one instance" in content:
                    print("\n[Error] IDEA is running. Please close it first.")
                    process.kill()
                    sys.exit(1)

        if process.poll() is None:
            process.kill()

    except Exception as e:
        print(f"[Error] IDEA execution failed: {e}")


def apply_patch_logic(original_file: str, formatted_file: str, vcs: VCS) -> bool:
    user_modified_lines = vcs.get_modified_lines(original_file)
    if not user_modified_lines:
        return True  # 没有修改行，视为成功（跳过）

    # 生成 Diff (使用 Git diff 作为通用的 Diff 引擎)
    cmd = ["git", "diff", "--no-index", "--unified=0", original_file, formatted_file]
    proc = subprocess.run(
        cmd, stdout=subprocess.PIPE, stderr=subprocess.PIPE, text=True, encoding="utf-8"
    )
    diff_output = proc.stdout

    if not diff_output:
        return True  # 无变化，视为成功

    filtered_hunks = []
    hunk_header_regex = re.compile(r"^@@ -(\d+)(?:,(\d+))? \+\d+(?:,\d+)? @@")

    current_hunk = []
    keep_current_hunk = False

    lines = diff_output.splitlines()
    start_idx = next((i for i, line in enumerate(lines) if line.startswith("@@")), 0)

    for line in lines[start_idx:]:
        if line.startswith("@@"):
            if current_hunk and keep_current_hunk:
                filtered_hunks.extend(current_hunk)

            current_hunk = [line]
            keep_current_hunk = False

            match = hunk_header_regex.match(line)
            if match:
                # 注意：这里我们取的是 '-' 侧 (original_file) 的行号
                # 因为我们要判断格式化工具是否动了 "用户修改过的旧行"
                start = int(match.group(1))
                length = int(match.group(2)) if match.group(2) else 1
                end = start + length - 1 if length > 0 else start

                # 只要存在交集，保留此 hunk
                if any(
                    line_num in user_modified_lines
                    for line_num in range(start, end + 1)
                ):
                    keep_current_hunk = True
        else:
            current_hunk.append(line)

    if current_hunk and keep_current_hunk:
        filtered_hunks.extend(current_hunk)

    if not filtered_hunks:
        return True  # 交集为空，丢弃修改

    # 手动构造 Patch Header
    git_path = original_file.replace("\\", "/")
    patch_header = [
        f"diff --git a/{git_path} b/{git_path}",
        f"--- a/{git_path}",
        f"+++ b/{git_path}",
    ]
    patch_content = "\n".join(patch_header + filtered_hunks) + "\n"

    # 应用 Patch (复用 git apply 作为优秀的补丁应用工具)
    p = subprocess.Popen(
        [
            "git",
            "apply",
            "--unidiff-zero",
            "--ignore-space-change",
            "--ignore-whitespace",
        ],
        stdin=subprocess.PIPE,
        stderr=subprocess.PIPE,
        text=True,
        encoding="utf-8",
    )
    out, err = p.communicate(input=patch_content)

    if p.returncode != 0:
        print(f"  [Error] Failed to apply patch for {original_file}")
        print(f"  [Git Apply Error] {err.strip()}")
        return False

    return True


# ================= 主程序入口 =================
def main():
    print("[Step 1] Detecting Version Control System...")
    vcs = detect_vcs()
    print(f"  -> Detected: {vcs.name}")

    # 自动切换到 Repo Root 确保相对路径无论在哪个子目录运行都不出错
    repo_root = vcs.get_repo_root()
    os.chdir(repo_root)

    repo_dir = os.path.basename(repo_root)
    print(f"[Step 2] Checking repository bounds (Repo: {repo_dir})")
    if repo_dir not in ALLOWED_REPOS:
        print(f"  -> Repo '{repo_dir}' not in allowed list, skipping.")
        sys.exit(0)

    print(f"[Step 3] Checking style XML at: {STYLE_XML_PATH}")
    if not os.path.exists(STYLE_XML_PATH):
        print("  -> [Error] Style XML not found!")
        sys.exit(1)

    print("[Step 4] Querying modified Java files...")
    files = vcs.get_modified_java_files()

    if not files:
        print("  -> No modified Java files to format. Exiting.")
        sys.exit(0)

    print(f"  -> Found {len(files)} modified Java files to format:")
    for f in files:
        print(f"     - {f}")

    idea_bin = get_platform_idea_bin()
    print(f"[Step 5] IDEA Formatter binary found at: {idea_bin}")

    with tempfile.TemporaryDirectory() as temp_root:
        print(f"[Step 6] Preparing temporary workspace at {temp_root} ...")

        temp_file_map = {}
        for f in files:
            dest_path = os.path.join(temp_root, f)
            os.makedirs(os.path.dirname(dest_path), exist_ok=True)
            shutil.copyfile(f, dest_path)
            temp_file_map[f] = dest_path

        print("[Step 7] Running IDEA Formatter (Batch Mode)...")
        run_idea_formatter_batch(idea_bin, temp_root)

        print("[Step 8] Processing diffs and applying partial formatting...")
        formatted_count = 0
        has_error = False

        for original_file, temp_file in temp_file_map.items():
            print(f"  -> Merging format diff for: {original_file}")
            success = apply_patch_logic(original_file, temp_file, vcs)
            if success:
                vcs.post_format_action(original_file)
                formatted_count += 1
            else:
                has_error = True

        if has_error:
            print("\n[Error] Formatting failed for some files. Aborting process.")
            sys.exit(1)

        print(f"\n[Success] Done. {formatted_count} files processed successfully.")


if __name__ == "__main__":
    main()