utils/find-dupes/find-dupes.py

#!/usr/bin/env python3
import argparse
import hashlib
import os
import sys
from collections import defaultdict
from pathlib import Path

CHUNK = 1024

G = "\033[32m"
Y = "\033[33m"
C = "\033[36m"
R = "\033[31m"
B = "\033[1m"
Z = "\033[0m"


def parse_size(s: str) -> int:
    s = s.strip().lower()
    mul = 1
    for suf, m in (
        ("kb", 1024), ("k", 1024),
        ("mb", 1024**2), ("m", 1024**2),
        ("gb", 1024**3), ("g", 1024**3),
        ("b", 1),
    ):
        if s.endswith(suf):
            s = s[:-len(suf)].strip()
            mul = m
            break
    return int(float(s) * mul)


def human(n: int) -> str:
    for unit in ("B", "KiB", "MiB", "GiB", "TiB"):
        if n < 1024 or unit == "TiB":
            return f"{n:.0f}{unit}" if unit == "B" else f"{n:.1f}{unit}"
        n /= 1024


def prog(stage: str, cur: int, total: int) -> None:
    msg = f"{stage:<9} | {cur}/{total} files"
    if not hasattr(prog, "_last"):
        prog._last = 0
    pad = max(0, prog._last - len(msg))
    sys.stderr.write("\r" + C + msg + Z + (" " * pad))
    sys.stderr.flush()
    prog._last = len(msg)


def parse_args():
    ap = argparse.ArgumentParser()
    ap.add_argument("root", nargs="?", default=".")
    ap.add_argument("--min-size", default="1")
    ap.add_argument("--max-size", default=None)
    ap.add_argument("--ignore", action="append", default=[])
    return ap.parse_args()


def is_ignored(p: Path, ignore: list[Path]) -> bool:
    try:
        rp = p.resolve()
    except OSError:
        return False
    return any(rp == q or q in rp.parents for q in ignore)


def file_head_tail(p: Path, n: int = CHUNK):
    try:
        size = p.stat().st_size
        with p.open("rb") as f:
            head = f.read(n)
            f.seek(max(0, size - n))
            tail = f.read(n)
        return head, tail
    except OSError:
        return None, None


def sha256sum(p: Path) -> str | None:
    h = hashlib.sha256()
    try:
        with p.open("rb") as f:
            while True:
                b = f.read(1024 * 1024)
                if not b:
                    break
                h.update(b)
        return h.hexdigest()
    except OSError:
        return None


def iter_files(root: Path, ignore: list[Path], min_size: int, max_size: int | None):
    for dirpath, dirnames, filenames in os.walk(root):
        dp = Path(dirpath)
        dirnames[:] = [d for d in dirnames if not (dp / d).is_symlink() and not is_ignored(dp / d, ignore)]

        for name in filenames:
            p = dp / name
            if p.is_symlink() or is_ignored(p, ignore):
                continue
            try:
                size = p.stat().st_size
            except OSError:
                continue
            if size < min_size:
                continue
            if max_size is not None and size > max_size:
                continue
            yield p, size


def main():
    args = parse_args()
    root = Path(args.root).expanduser().resolve()
    min_size = parse_size(args.min_size)
    max_size = parse_size(args.max_size) if args.max_size else None
    ignore = [Path(p).expanduser().resolve() for p in args.ignore]

    by_size = defaultdict(list)
    total = 0

    for p, size in iter_files(root, ignore, min_size, max_size):
        total += 1
        if total % 100 == 0:
            prog("Scanning", total, total)
        by_size[size].append(p)

    if not total:
        print("No files found.")
        return

    prog("Scanning", total, total)
    sys.stderr.write("\n")

    candidates = {s: ps for s, ps in by_size.items() if len(ps) > 1}
    if not candidates:
        print("No duplicates found.")
        return

    by_ht = defaultdict(list)
    ht_total = sum(len(ps) for ps in candidates.values())
    done = 0
    for size, ps in candidates.items():
        for p in ps:
            done += 1
            prog("Head/tail", done, ht_total)
            h, t = file_head_tail(p)
            if h is not None:
                by_ht[(size, h, t)].append(p)
    sys.stderr.write("\n")

    ht_groups = {k: ps for k, ps in by_ht.items() if len(ps) > 1}
    if not ht_groups:
        print("No duplicates found.")
        return

    by_hash = defaultdict(list)
    hash_total = sum(len(ps) for ps in ht_groups.values())
    done = 0
    for ps in ht_groups.values():
        for p in ps:
            done += 1
            prog("Hashing", done, hash_total)
            h = sha256sum(p)
            if h:
                by_hash[h].append(p)
    sys.stderr.write("\n")

    groups = []
    for h, ps in by_hash.items():
        if len(ps) > 1:
            try:
                size = ps[0].stat().st_size
            except OSError:
                size = 0
            groups.append((len(ps), size, h, ps))

    if not groups:
        print("No duplicates found.")
        return

    groups.sort(key=lambda x: (-x[0], -x[1]))

    for idx, (count, size, h, ps) in enumerate(groups, 1):
        print(f"{B}{Y}Group {idx}{Z} ({count} files, {human(size)}) {C}{h[:12]}{Z}")
        for p in sorted(ps):
            print(f"  {G}{p.resolve()}{Z}")


if __name__ == "__main__":
    main()