dupfinder.py

#!/usr/bin/env python3
"""
dupfinder — 查找并清理重复文件

基于 SHA-256 哈希检测重复文件，支持交互式或自动清理。
重复文件按哈希分组，默认保留每组第一个文件（可配置保留策略）。

用法:
  # 仅扫描，列出重复文件
  python dupfinder.py /path/to/dir

  # 交互模式：每遇到一组重复文件时询问保留哪个
  python dupfinder.py /path/to/dir --interactive

  # 自动删除重复文件，保留每组中最早的文件
  python dupfinder.py /path/to/dir --delete --keep oldest

  # 预览模式：显示会删除哪些文件但不实际删除
  python dupfinder.py /path/to/dir --delete --dry-run

  # 排除指定目录
  python dupfinder.py /path/to/dir --exclude node_modules --exclude .git

  # 限制扫描的最小文件大小（跳过太小的文件）
  python dupfinder.py /path/to/dir --min-size 1K

保留策略 (--keep):
  first     保留路径字典序排第一的文件（默认）
  last      保留路径字典序排最后的文件
  oldest    保留修改时间最早的文件
  newest    保留修改时间最新的文件
  shortest  保留文件名字符数最少的
  longest   保留文件名字符数最多的
"""

import argparse
import hashlib
import os
import sys
from collections import defaultdict
from pathlib import Path
from typing import Dict, List, Set

# ── ANSI 颜色（终端友好输出） ──────────────────────────────────────────────
C_RESET = "\033[0m"
C_BOLD = "\033[1m"
C_DIM = "\033[2m"
C_RED = "\033[91m"
C_GREEN = "\033[92m"
C_YELLOW = "\033[93m"
C_CYAN = "\033[96m"


def parse_size(size_str: str) -> int:
    """解析带单位的大小字符串，如 1K, 2M, 500B。"""
    multipliers = {"B": 1, "K": 1024, "KB": 1024,
                   "M": 1024**2, "MB": 1024**2,
                   "G": 1024**3, "GB": 1024**3}
    size_str = size_str.strip().upper()
    for suffix, mult in multipliers.items():
        if size_str.endswith(suffix):
            try:
                return int(float(size_str[:-len(suffix)]) * mult)
            except ValueError:
                pass
    try:
        return int(size_str)
    except ValueError:
        raise argparse.ArgumentTypeError(f"无法解析大小: {size_str}")


def format_size(num_bytes: int) -> str:
    """以人类可读格式显示字节数。"""
    for unit in ("B", "KB", "MB", "GB", "TB"):
        if num_bytes < 1024:
            return f"{num_bytes:.1f} {unit}"
        num_bytes /= 1024
    return f"{num_bytes:.1f} PB"


def hash_file(filepath: Path, chunk_size: int = 65536) -> str:
    """返回文件的 SHA-256 十六进制哈希值。"""
    sha256 = hashlib.sha256()
    with open(filepath, "rb") as f:
        while True:
            chunk = f.read(chunk_size)
            if not chunk:
                break
            sha256.update(chunk)
    return sha256.hexdigest()


def collect_files(root: Path, excludes: Set[str], min_size: int) -> Dict[int, List[Path]]:
    """
    遍历目录，将文件分桶到 {文件大小: [路径列表]}。
    用文件大小作为第一轮过滤，避免对每个文件都计算哈希。
    """
    by_size: Dict[int, List[Path]] = defaultdict(list)
    total = 0

    for dirpath, dirnames, filenames in os.walk(root, followlinks=False):
        # 规范化路径用于排除匹配
        rel = os.path.relpath(dirpath, root)
        dirnames[:] = [d for d in dirnames
                       if os.path.join(rel, d) not in excludes]

        for fname in filenames:
            full = Path(dirpath) / fname
            try:
                stat = full.lstat()
            except OSError as e:
                print(f"{C_DIM}⚠ 跳过（无法读取）: {full} ({e}){C_RESET}")
                continue

            if not stat.st_mode & 0o100000:  # S_IFREG — 只处理普通文件
                continue
            if stat.st_size < min_size:
                continue

            by_size[stat.st_size].append(full)
            total += 1
            if total % 1000 == 0:
                print(f"{C_DIM}  已扫描 {total} 个文件...{C_RESET}", file=sys.stderr)

    return by_size


def find_duplicates(root: Path,
                    excludes: Set[str],
                    min_size: int) -> Dict[str, List[Path]]:
    """
    返回 {哈希值: [相同文件的路径列表]}，
    只对大小桶中多于 1 个文件的桶计算哈希，大幅减少计算量。
    """
    print(f"{C_CYAN}🔍 扫描目录: {root}{C_RESET}")
    by_size = collect_files(root, excludes, min_size)

    # 只对大小冲突的组计算哈希
    candidates = [(size, paths) for size, paths in by_size.items()
                  if len(paths) > 1]
    print(f"{C_DIM}  大小冲突组: {len(candidates)}（共 {sum(len(p) for _, p in candidates)} 个待哈希文件）{C_RESET}")

    by_hash: Dict[str, List[Path]] = defaultdict(list)
    hashed = 0
    for size, paths in candidates:
        for p in paths:
            try:
                h = hash_file(p)
                by_hash[h].append(p)
            except OSError as e:
                print(f"{C_DIM}⚠ 无法读取: {p} ({e}){C_RESET}")
            hashed += 1
            if hashed % 100 == 0:
                print(f"{C_DIM}  哈希进度: {hashed}/{sum(len(p) for _, p in candidates)}{C_RESET}", file=sys.stderr)

    # 只返回真正的重复组
    return {h: paths for h, paths in by_hash.items() if len(paths) > 1}


def pick_kept(dupes: List[Path], strategy: str) -> Path:
    """根据策略选择保留哪个文件，返回被保留的路径。"""
    if strategy == "last":
        return sorted(dupes)[-1]
    if strategy == "oldest":
        return min(dupes, key=lambda p: p.stat().st_mtime)
    if strategy == "newest":
        return max(dupes, key=lambda p: p.stat().st_mtime)
    if strategy == "shortest":
        return min(dupes, key=lambda p: len(p.name))
    if strategy == "longest":
        return max(dupes, key=lambda p: len(p.name))
    # "first" — 按路径排序的第一个
    return sorted(dupes)[0]


def interactive_pick(dupes: List[Path], file_hash: str = "") -> Path:
    """交互式选择保留哪个文件。"""
    print(f"\n{C_YELLOW}{C_BOLD}━━━ 重复文件组（SHA-256: {file_hash}）━━━{C_RESET}")
    for i, p in enumerate(dupes):
        try:
            mtime = p.stat().st_mtime
            from datetime import datetime
            ts = datetime.fromtimestamp(mtime).strftime("%Y-%m-%d %H:%M")
        except OSError:
            ts = "未知"
        size = format_size(p.stat().st_size if p.exists() else 0)
        print(f"  {C_GREEN}[{i}] {p}{C_RESET}  {C_DIM}({size}, {ts}){C_RESET}")

    while True:
        try:
            choice = input(f"{C_CYAN}保留哪个？ [0-{len(dupes)-1}, s=跳过, q=退出]: {C_RESET}").strip()
        except (EOFError, KeyboardInterrupt):
            print()
            sys.exit(0)

        if choice.lower() == "q":
            print("已退出。")
            sys.exit(0)
        if choice.lower() == "s":
            return None  # type: ignore — 跳过此组
        try:
            idx = int(choice)
            if 0 <= idx < len(dupes):
                return dupes[idx]
        except ValueError:
            pass
        print(f"{C_RED}无效选择，请重试。{C_RESET}")


def main():
    parser = argparse.ArgumentParser(
        description="基于 SHA-256 的重复文件查找与清理工具",
        formatter_class=argparse.RawDescriptionHelpFormatter,
        epilog=__doc__,
    )
    parser.add_argument("path", type=Path,
                        help="要扫描的根目录路径")
    parser.add_argument("--interactive", "-i", action="store_true",
                        help="交互模式：每组重复文件时询问保留哪个")
    parser.add_argument("--delete", "-d", action="store_true",
                        help="删除重复文件（配合 --keep 指定保留策略）")
    parser.add_argument("--keep", "-k", type=str,
                        choices=["first", "last", "oldest", "newest", "shortest", "longest"],
                        default="first",
                        help="保留策略 (默认: first)")
    parser.add_argument("--dry-run", "-n", action="store_true",
                        help="预览模式：仅显示将要做什么，不实际删除")
    parser.add_argument("--min-size", "-s", type=parse_size, default="1",
                        help="最小文件大小，支持单位 B/K/M/G (默认: 1 字节)")
    parser.add_argument("--exclude", "-e", action="append", default=[],
                        dest="excludes",
                        help="排除的子目录名（可重复使用），相对于扫描根目录")
    parser.add_argument("--yes", "-y", action="store_true",
                        help="跳过所有确认提示（非交互式删除时生效）")
    args = parser.parse_args()

    root = args.path.resolve()
    if not root.is_dir():
        print(f"{C_RED}错误: 路径不存在或不是目录: {root}{C_RESET}")
        sys.exit(1)

    excludes: Set[str] = set(args.excludes)

    # ── 查找重复文件 ──
    dupes_map = find_duplicates(root, excludes, args.min_size)

    if not dupes_map:
        print(f"{C_GREEN}✅ 未发现重复文件。{C_RESET}")
        sys.exit(0)

    total_dup_groups = len(dupes_map)
    total_dup_files = sum(len(v) for v in dupes_map.values())
    total_wasted = 0  # 可节省的空间

    print(f"\n{C_YELLOW}{C_BOLD}发现 {total_dup_groups} 组重复文件（共 {total_dup_files} 个文件）：{C_RESET}\n")

    for idx, (h, files) in enumerate(dupes_map.items(), 1):
        size = format_size(files[0].stat().st_size)
        wasted = (len(files) - 1) * files[0].stat().st_size
        total_wasted += wasted

        print(f"{C_BOLD}组 {idx}/{total_dup_groups}  —  {size} 每个  (SHA-256: {h}){C_RESET}")
        for p in sorted(files):
            print(f"    {p}")
        print()

    print(f"{C_YELLOW}📊 汇总: {total_dup_groups} 组重复, {total_dup_files} 个文件, "
          f"可释放 {format_size(total_wasted)}{C_RESET}")

    # ── 处理（交互式 / 自动删除 / 仅报告） ──
    if args.interactive:
        # 交互模式
        print(f"\n{C_CYAN}进入交互模式，逐个处理每组重复文件...{C_RESET}")
        to_delete: List[Path] = []
        skipped = 0
        for h, files in dupes_map.items():
            kept = interactive_pick(files, h)
            if kept is None:
                skipped += 1
                continue
            to_delete.extend(p for p in files if p != kept)

        if not to_delete:
            print(f"{C_GREEN}没有文件需要删除。{C_RESET}")
            sys.exit(0)

        print(f"\n{C_RED}{C_BOLD}即将删除 {len(to_delete)} 个文件，共 {format_size(sum(p.stat().st_size for p in to_delete))}{C_RESET}")
        ans = input(f"{C_RED}确认删除？[y/N]: {C_RESET}").strip().lower()
        if ans not in ("y", "yes"):
            print("取消。")
            sys.exit(0)

        for p in to_delete:
            try:
                p.unlink()
                print(f"{C_RED}  已删除: {p}{C_RESET}")
            except OSError as e:
                print(f"{C_RED}  删除失败: {p} ({e}){C_RESET}")
        print(f"{C_GREEN}✅ 完成！释放 {format_size(total_wasted - skipped * files[0].stat().st_size if dupes_map else 0)}{C_RESET}")

    elif args.delete:
        # 自动删除模式
        to_delete: List[Path] = []
        for h, files in dupes_map.items():
            kept = pick_kept(files, args.keep)
            to_delete.extend(p for p in files if p != kept)

        total_delete_size = sum(p.stat().st_size for p in to_delete)

        print(f"\n{C_RED}保留策略: {C_BOLD}{args.keep}{C_RESET}")
        print(f"{C_RED}将删除 {len(to_delete)} 个重复文件，释放 {format_size(total_delete_size)}{C_RESET}")

        if args.dry_run:
            print(f"\n{C_CYAN}── 预览（不会实际删除）──{C_RESET}")
            for h, files in dupes_map.items():
                kept = pick_kept(files, args.keep)
                print(f"\n{C_BOLD}保留:{C_RESET} {C_GREEN}{kept}{C_RESET}")
                for p in sorted(files):
                    if p != kept:
                        print(f"  {C_RED}删除:{C_RESET} {p}")
            print(f"\n{C_DIM}(dry-run 模式，未做任何更改){C_RESET}")
        else:
            if not args.yes:
                ans = input(f"\n{C_RED}确认删除？[y/N]: {C_RESET}").strip().lower()
                if ans not in ("y", "yes"):
                    print("取消。")
                    sys.exit(0)

            deleted = 0
            failed = 0
            for p in to_delete:
                try:
                    p.unlink()
                    print(f"{C_RED}  已删除: {p}{C_RESET}")
                    deleted += 1
                except OSError as e:
                    print(f"{C_RED}  删除失败: {p} ({e}){C_RESET}")
                    failed += 1
            print(f"\n{C_GREEN}✅ 完成: {deleted} 个文件已删除, {failed} 个失败, 释放 {format_size(total_delete_size)}{C_RESET}")

    else:
        # 仅报告模式
        print(f"\n{C_DIM}提示: 使用 --delete 自动删除, --interactive 逐个确认, "
              f"或 --delete --dry-run 预览{C_RESET}")


if __name__ == "__main__":
    main()