zip_source.py

#!/usr/bin/env python3
"""
zip_source.py — Zip a codebase's source files, intelligently skipping
dev/build/generated artifacts (node_modules, dist, __pycache__, etc.).

Usage:
    python zip_source.py
    python zip_source.py --path /path/to/project --dest /path/to/output
"""

import os
import sys
import zipfile
import argparse
import datetime
import platform
import unicodedata
from pathlib import Path

# ---------------------------------------------------------------------------
# Exclusion rules
# ---------------------------------------------------------------------------

# Directories that are NEVER included (matched against the directory basename).
EXCLUDED_DIRS: set[str] = {
    # JavaScript / Node
    "node_modules", ".npm", ".yarn", ".pnp",
    # Build outputs
    "dist", "build", "out", "output", "outputs",
    ".next", ".nuxt", ".svelte-kit", ".output",
    "public/build",                    # common SPA convention
    # Python
    "__pycache__", ".venv", "venv", "env", ".env",
    "site-packages", "dist-info", "egg-info",
    ".eggs", "*.egg-info",
    ".tox", ".nox",
    # Java / Kotlin / Scala
    "target", ".gradle", ".mvn",
    # Ruby
    "vendor/bundle", ".bundle",
    # Rust
    "target",                          # also Rust
    # Go
    "vendor",
    # .NET
    "bin", "obj", "packages",
    # IDEs & editors
    ".idea", ".vscode", ".vs",
    ".eclipse", ".settings",
    # Version control
    ".git", ".svn", ".hg",
    # Coverage & testing artefacts
    "coverage", ".coverage", "htmlcov",
    ".pytest_cache", ".mypy_cache",
    ".hypothesis", ".ruff_cache",
    # Caches & temp
    ".cache", ".parcel-cache", ".turbo",
    ".webpack", ".rollup.cache",
    "tmp", "temp", ".tmp",
    # Terraform / infra
    ".terraform",
    # macOS
    ".DS_Store",                       # file, but catch as dir too
}

# File extensions that are excluded (generated, compiled, binary, etc.)
EXCLUDED_EXTENSIONS: set[str] = {
    # Compiled / binary
    ".pyc", ".pyo", ".pyd",
    ".class", ".jar", ".war", ".ear",
    ".o", ".obj", ".a", ".lib", ".so", ".dylib", ".dll", ".exe",
    ".wasm",
    # Compiled web
    ".map",                            # source maps — generated
    # Archives (don't nest zips)
    ".zip", ".tar", ".gz", ".tgz", ".bz2", ".xz", ".rar", ".7z",
    # Heavyweight design assets (not source code)
    ".psd", ".ai", ".sketch", ".fig",
    # DB / data dumps
    ".sqlite", ".sqlite3", ".db",
    ".dump", ".sql",            # SQL dumps are usually large; keep .sql schema files? debatable
    # Logs
    ".log",
    # OS artefacts
    ".DS_Store", ".Thumbs.db",
}

# Exact filenames that are always excluded.
EXCLUDED_FILENAMES: set[str] = {
    ".DS_Store",
    "Thumbs.db",
    "desktop.ini",
    "npm-debug.log",
    "yarn-error.log",
    "yarn-debug.log",
    ".env",                            # secrets — never zip
    ".env.local",
    ".env.production",
    ".env.development",
    ".env.test",
    # Lock files (large, regeneratable)
    "package-lock.json",
    "yarn.lock",
    "pnpm-lock.yaml",
    "Pipfile.lock",
    "poetry.lock",
    "Gemfile.lock",
    "Cargo.lock",
    "composer.lock",
    "go.sum",
}

# ---------------------------------------------------------------------------
# Media toggle
# ---------------------------------------------------------------------------

# Set to True to include images, video, and audio files in the zip.
# Useful when your project has a public/ or assets/ folder with media
# that is part of the source (e.g. static sites, mobile apps).
# Set to False to exclude all media files (default: False).
INCLUDE_MEDIA: bool = False

MEDIA_EXTENSIONS: set[str] = {
    # Images
    ".png", ".jpg", ".jpeg", ".gif", ".webp", ".svg",
    ".ico", ".bmp", ".tiff", ".tif", ".avif", ".heic", ".heif",
    # Video
    ".mp4", ".mov", ".avi", ".mkv", ".webm", ".flv", ".wmv", ".m4v",
    # Audio
    ".mp3", ".wav", ".flac", ".ogg", ".aac", ".m4a", ".wma",
}

# File size limit per individual file (default 50 MB). Files larger than this
# are skipped and reported. Protects against accidentally zipping huge assets.
MAX_FILE_BYTES: int = 50 * 1024 * 1024   # 50 MB


# ---------------------------------------------------------------------------
# Helpers
# ---------------------------------------------------------------------------

def sanitize_zip_path(path: str) -> str:
    """
    Normalise a path string for safe inclusion inside a ZIP archive:
    - Replace backslashes with forward slashes.
    - Strip leading slashes / drive letters.
    - Normalise unicode to NFC (avoids macOS vs Linux decomposition mismatches).
    """
    path = unicodedata.normalize("NFC", path)
    path = path.replace("\\", "/")
    # Remove dangerous leading components
    while path.startswith(("/", "../", "./")):
        path = path.lstrip("/").lstrip(".")
        if path.startswith("/"):
            path = path[1:]
    return path


def is_excluded_dir(dirpath: Path) -> bool:
    """Return True if this directory should be skipped entirely."""
    name = dirpath.name
    # Direct name match
    if name in EXCLUDED_DIRS:
        return True
    # Glob-style suffix match (e.g. "*.egg-info")
    for pattern in EXCLUDED_DIRS:
        if pattern.startswith("*") and name.endswith(pattern[1:]):
            return True
    return False


def is_excluded_file(filepath: Path) -> bool:
    """Return True if this file should be skipped."""
    if filepath.name in EXCLUDED_FILENAMES:
        return True
    ext = filepath.suffix.lower()
    if not INCLUDE_MEDIA and ext in MEDIA_EXTENSIONS:
        return True
    if ext in EXCLUDED_EXTENSIONS:
        return True
    return False


def human_readable_size(num_bytes: int) -> str:
    for unit in ("B", "KB", "MB", "GB"):
        if num_bytes < 1024:
            return f"{num_bytes:.1f} {unit}"
        num_bytes /= 1024
    return f"{num_bytes:.1f} TB"


def collect_files(root: Path) -> tuple[list[Path], list[str], list[str]]:
    """
    Walk *root* and return:
        included  — list of absolute Path objects to zip
        skipped   — human-readable messages for skipped files
        skipped_dirs — human-readable messages for skipped directory trees
    """
    included: list[Path] = []
    skipped: list[str] = []
    skipped_dirs: list[str] = []

    for dirpath, dirnames, filenames in os.walk(root, topdown=True, followlinks=False):
        current_dir = Path(dirpath)

        # Prune excluded dirs in-place so os.walk won't descend into them.
        pruned = []
        for d in list(dirnames):
            sub = current_dir / d
            if is_excluded_dir(sub):
                skipped_dirs.append(str(sub.relative_to(root)))
                pruned.append(d)
        for d in pruned:
            dirnames.remove(d)

        for fname in filenames:
            fpath = current_dir / fname

            # Skip broken symlinks
            if fpath.is_symlink() and not fpath.exists():
                skipped.append(f"  [broken symlink] {fpath.relative_to(root)}")
                continue

            if is_excluded_file(fpath):
                skipped.append(f"  [excluded type]  {fpath.relative_to(root)}")
                continue

            try:
                size = fpath.stat().st_size
            except OSError as exc:
                skipped.append(f"  [stat error]     {fpath.relative_to(root)}: {exc}")
                continue

            is_media = fpath.suffix.lower() in MEDIA_EXTENSIONS
            if size > MAX_FILE_BYTES and not (INCLUDE_MEDIA and is_media):
                skipped.append(
                    f"  [too large]      {fpath.relative_to(root)} "
                    f"({human_readable_size(size)})"
                )
                continue

            included.append(fpath)

    return included, skipped, skipped_dirs


def prompt_path(prompt_text: str, must_exist: bool = False) -> Path:
    """Prompt the user for a filesystem path, with basic validation."""
    while True:
        raw = input(prompt_text).strip()
        if not raw:
            print("  ⚠  No input provided. Please try again.")
            continue

        # Expand ~ and env vars
        expanded = Path(os.path.expandvars(os.path.expanduser(raw)))

        if must_exist and not expanded.exists():
            print(f"  ⚠  Path does not exist: {expanded}")
            continue
        if must_exist and not expanded.is_dir():
            print(f"  ⚠  Path is not a directory: {expanded}")
            continue

        return expanded


def resolve_destination(source: Path, dest_input: str) -> Path:
    """
    Resolve where the zip file will be written.

    Rules
    -----
    * If dest_input is empty → parent of source directory.
    * If dest_input is a directory (existing or ending with /) → place zip inside it.
    * If dest_input ends with .zip → use it literally as the zip path.
    * Otherwise → treat as a directory, place zip inside it.
    """
    timestamp = datetime.datetime.now().strftime("%Y%m%d_%H%M%S")
    zip_name = f"{source.name}_{timestamp}.zip"

    if not dest_input.strip():
        return source.parent / zip_name

    dest = Path(os.path.expandvars(os.path.expanduser(dest_input.strip())))

    if dest_input.strip().endswith(".zip"):
        # User specified an explicit zip filename
        return dest

    # Treat as directory
    return dest / zip_name


def ensure_parent_dir(zip_path: Path) -> None:
    """Create parent directories for the zip file if they don't exist."""
    try:
        zip_path.parent.mkdir(parents=True, exist_ok=True)
    except OSError as exc:
        sys.exit(f"✗  Cannot create destination directory {zip_path.parent}: {exc}")


def check_no_overwrite(zip_path: Path) -> Path:
    """
    If zip_path already exists, append an incrementing counter to avoid
    silently overwriting an existing archive.
    """
    if not zip_path.exists():
        return zip_path
    stem = zip_path.stem
    suffix = zip_path.suffix
    parent = zip_path.parent
    counter = 1
    while True:
        candidate = parent / f"{stem}_{counter}{suffix}"
        if not candidate.exists():
            print(f"  ℹ  Destination exists. Writing to: {candidate.name}")
            return candidate
        counter += 1


def write_zip(
    zip_path: Path,
    source_root: Path,
    files: list[Path],
) -> tuple[int, int]:
    """
    Write files into the zip archive.

    Returns (files_written, bytes_written).
    """
    files_written = 0
    bytes_written = 0

    with zipfile.ZipFile(zip_path, "w", compression=zipfile.ZIP_DEFLATED, allowZip64=True) as zf:
        for fpath in files:
            arcname = sanitize_zip_path(str(fpath.relative_to(source_root)))
            try:
                zf.write(fpath, arcname=arcname)
                files_written += 1
                bytes_written += fpath.stat().st_size
            except PermissionError:
                print(f"  ⚠  Permission denied, skipping: {fpath.relative_to(source_root)}")
            except OSError as exc:
                print(f"  ⚠  OS error for {fpath.relative_to(source_root)}: {exc}")

    return files_written, bytes_written


# ---------------------------------------------------------------------------
# CLI / interactive entry-point
# ---------------------------------------------------------------------------

def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser(
        description="Zip a codebase's source files, skipping dev/build artifacts.",
        formatter_class=argparse.RawDescriptionHelpFormatter,
        epilog=__doc__,
    )
    parser.add_argument("--path", "-p", help="Path to the codebase directory.")
    parser.add_argument("--dest", "-d", help="Destination path or directory for the zip file.")
    parser.add_argument(
        "--verbose", "-v", action="store_true",
        help="Print every skipped file (default: summary only).",
    )
    parser.add_argument(
        "--max-file-mb", type=float, default=50,
        help="Per-file size limit in MB (default: 50). Files exceeding this are skipped.",
    )
    return parser.parse_args()


def main() -> None:
    args = parse_args()

    # Override global limit if user passed --max-file-mb
    global MAX_FILE_BYTES
    MAX_FILE_BYTES = int(args.max_file_mb * 1024 * 1024)

    print()
    print("═" * 60)
    print("  zip_source — Smart Codebase Zipper")
    print("  Skips node_modules, dist, build artefacts, secrets, …")
    print("═" * 60)
    print()

    # ── 1. Source path ──────────────────────────────────────────
    if args.path:
        source = Path(os.path.expandvars(os.path.expanduser(args.path)))
        if not source.exists():
            sys.exit(f"✗  Provided path does not exist: {source}")
        if not source.is_dir():
            sys.exit(f"✗  Provided path is not a directory: {source}")
    else:
        print("Step 1 of 2 — Source directory")
        source = prompt_path("  Enter the path to your codebase folder: ", must_exist=True)

    source = source.resolve()
    print(f"\n  ✔  Source  : {source}")

    # ── 2. Destination ──────────────────────────────────────────
    if args.dest:
        dest_input = args.dest
    else:
        print("\nStep 2 of 2 — Destination")
        print(f"  Default (press Enter): {source.parent}")
        dest_input = input("  Enter destination path (or press Enter for default): ").strip()

    zip_path = resolve_destination(source, dest_input)
    ensure_parent_dir(zip_path)
    zip_path = zip_path.resolve()

    # Guard: don't zip the destination into itself
    try:
        zip_path.relative_to(source)
        sys.exit(
            "✗  Destination is inside the source directory. "
            "Please choose a different destination."
        )
    except ValueError:
        pass  # Good — destination is outside source

    zip_path = check_no_overwrite(zip_path)
    print(f"  ✔  Output  : {zip_path}")

    # ── 3. Collect files ────────────────────────────────────────
    print("\n  Scanning codebase …")
    included, skipped_files, skipped_dirs = collect_files(source)

    total_found = len(included) + len(skipped_files)
    print(f"  Found   : {total_found:,} candidate files")
    print(f"  Skipped : {len(skipped_files):,} files  |  {len(skipped_dirs):,} directory trees")
    print(f"  To zip  : {len(included):,} files")

    if not included:
        sys.exit("\n✗  No source files found to zip. Nothing written.")

    # ── 4. Optionally show exclusion details ────────────────────
    has_exclusions = skipped_files or skipped_dirs
    if has_exclusions:
        print()
        show = input("  Show excluded items before proceeding? [y/N]: ").strip().lower()
        if show in ("y", "yes") or args.verbose:
            if skipped_dirs:
                print()
                print("  ── Excluded directory trees " + "─" * 31)
                for d in skipped_dirs:
                    print(f"    [dir]  {d}")
            if skipped_files:
                print()
                print("  ── Excluded files " + "─" * 40)
                for msg in skipped_files:
                    print(f"    {msg.strip()}")
            print()
            dir_word = "directory" if len(skipped_dirs) == 1 else "directories"
            file_word = "file" if len(skipped_files) == 1 else "files"
            print(f"  ({len(skipped_dirs)} {dir_word}, {len(skipped_files)} {file_word} excluded)")

    # ── 5. Confirm ──────────────────────────────────────────────
    print()
    confirm = input("  Proceed? [Y/n]: ").strip().lower()
    if confirm not in ("", "y", "yes"):
        sys.exit("  Aborted.")

    # ── 5. Write zip ────────────────────────────────────────────
    print("\n  Compressing …")
    try:
        files_written, bytes_written = write_zip(zip_path, source, included)
    except KeyboardInterrupt:
        # Clean up partial zip
        if zip_path.exists():
            zip_path.unlink()
        sys.exit("\n\n  Interrupted. Partial zip removed.")
    except OSError as exc:
        if zip_path.exists():
            zip_path.unlink()
        sys.exit(f"\n✗  Failed to write zip: {exc}")

    # ── 6. Summary ──────────────────────────────────────────────
    zip_size = zip_path.stat().st_size
    ratio = (1 - zip_size / bytes_written) * 100 if bytes_written else 0

    print()
    print("═" * 60)
    print("  ✅  Done!")
    print(f"  Files archived : {files_written:,}")
    print(f"  Source size    : {human_readable_size(bytes_written)}")
    print(f"  Zip size       : {human_readable_size(zip_size)}  ({ratio:.1f}% compression)")
    print(f"  Saved to       : {zip_path}")
    print("═" * 60)
    print()

    if platform.system() == "Darwin":
        print("  Tip: open the containing folder with:")
        print(f"       open \"{zip_path.parent}\"")
    elif platform.system() == "Windows":
        print("  Tip: open the containing folder with:")
        print(f"       explorer \"{zip_path.parent}\"")


if __name__ == "__main__":
    main()