scripts/clear_wjad_extract_state.py

"""删除 ``jobs_extract_archives.py`` 在 ``--out-root`` 下生成的 ``_wjad_extract_state`` 目录。

只删「已解压」进度标记，**不会**删除 ``extracted/`` 里真正的数据集文件。
删完后可重新跑解压脚本，会按归档从头再解压一遍（若目录已存在可能覆盖/混合，按需先清 extracted）。

用法示例::

    # 本机 / Bucket 挂载路径与解压时 --out-root 一致
    python scripts/clear_wjad_extract_state.py --out-root /mnt/wjad/extracted/cosmos_hub

    # 只看会删什么，不真删
    python scripts/clear_wjad_extract_state.py --out-root F:/bucket/extracted/cosmos_hub --dry-run

环境变量（可选）::

    设置 WJAD_EXTRACTED_ROOT 后可省略 --out-root

    Linux Job::

        export WJAD_EXTRACTED_ROOT=/mnt/wjad/extracted/cosmos_hub
        python scripts/clear_wjad_extract_state.py
"""

from __future__ import annotations

import argparse
import os
import shutil
import sys
from pathlib import Path

STATE_DIRNAME = "_wjad_extract_state"


def main() -> None:
    p = argparse.ArgumentParser(description="Remove _wjad_extract_state under extracted out-root.")
    p.add_argument(
        "--out-root",
        type=Path,
        default=None,
        help="与 jobs_extract_archives --out-root 相同；不设则用环境变量 WJAD_EXTRACTED_ROOT",
    )
    p.add_argument("--dry-run", action="store_true", help="只打印路径，不删除")
    args = p.parse_args()
    out = args.out_root
    if out is None:
        env = os.environ.get("WJAD_EXTRACTED_ROOT")
        if not env:
            print(
                "需要 --out-root 或环境变量 WJAD_EXTRACTED_ROOT",
                file=sys.stderr,
            )
            sys.exit(2)
        out = Path(env)
    out = out.resolve()
    state = out / STATE_DIRNAME
    if not state.exists():
        print(f"[clear] 不存在，跳过: {state}")
        return
    if not state.is_dir():
        print(f"[clear] 不是目录，拒绝: {state}", file=sys.stderr)
        sys.exit(1)
    print(f"[clear] {'(dry-run) ' if args.dry_run else ''}目标: {state}")
    if args.dry_run:
        n = sum(1 for _ in state.rglob("*"))
        print(f"[clear] dry-run: 其下约 {n} 个条目（含文件与目录）")
        return
    shutil.rmtree(state)
    print(f"[clear] 已删除: {state}")


if __name__ == "__main__":
    main()