codecount.py

#!/usr/bin/env python3
"""CodeCount - Source code line counter. Count lines of code, comments, and blanks by language."""

import os
import sys
import json
import argparse
import fnmatch
from pathlib import Path

LANGUAGES = {
    ".py":     {"name": "Python",       "line": "#",    "block": ('"""', '"""'), "alt_block": ("'''", "'''")},
    ".js":     {"name": "JavaScript",   "line": "//",   "block": ("/*", "*/")},
    ".jsx":    {"name": "JSX",          "line": "//",   "block": ("/*", "*/")},
    ".ts":     {"name": "TypeScript",   "line": "//",   "block": ("/*", "*/")},
    ".tsx":    {"name": "TSX",          "line": "//",   "block": ("/*", "*/")},
    ".java":   {"name": "Java",         "line": "//",   "block": ("/*", "*/")},
    ".c":      {"name": "C",            "line": "//",   "block": ("/*", "*/")},
    ".h":      {"name": "C Header",     "line": "//",   "block": ("/*", "*/")},
    ".cpp":    {"name": "C++",          "line": "//",   "block": ("/*", "*/")},
    ".hpp":    {"name": "C++ Header",   "line": "//",   "block": ("/*", "*/")},
    ".cs":     {"name": "C#",           "line": "//",   "block": ("/*", "*/")},
    ".go":     {"name": "Go",           "line": "//",   "block": ("/*", "*/")},
    ".rs":     {"name": "Rust",         "line": "//",   "block": ("/*", "*/")},
    ".rb":     {"name": "Ruby",         "line": "#",    "block": ("=begin", "=end")},
    ".php":    {"name": "PHP",          "line": "//",   "block": ("/*", "*/")},
    ".swift":  {"name": "Swift",        "line": "//",   "block": ("/*", "*/")},
    ".kt":     {"name": "Kotlin",       "line": "//",   "block": ("/*", "*/")},
    ".scala":  {"name": "Scala",        "line": "//",   "block": ("/*", "*/")},
    ".r":      {"name": "R",            "line": "#"},
    ".R":      {"name": "R",            "line": "#"},
    ".lua":    {"name": "Lua",          "line": "--",   "block": ("--[[", "]]")},
    ".pl":     {"name": "Perl",         "line": "#"},
    ".pm":     {"name": "Perl Module",  "line": "#"},
    ".sh":     {"name": "Shell",        "line": "#"},
    ".bash":   {"name": "Bash",         "line": "#"},
    ".zsh":    {"name": "Zsh",          "line": "#"},
    ".html":   {"name": "HTML",         "block": ("<!--", "-->")},
    ".htm":    {"name": "HTML",         "block": ("<!--", "-->")},
    ".css":    {"name": "CSS",          "block": ("/*", "*/")},
    ".scss":   {"name": "SCSS",         "line": "//",   "block": ("/*", "*/")},
    ".sass":   {"name": "Sass",         "line": "//"},
    ".less":   {"name": "Less",         "line": "//",   "block": ("/*", "*/")},
    ".sql":    {"name": "SQL",          "line": "--",   "block": ("/*", "*/")},
    ".yaml":   {"name": "YAML",         "line": "#"},
    ".yml":    {"name": "YAML",         "line": "#"},
    ".toml":   {"name": "TOML",         "line": "#"},
    ".xml":    {"name": "XML",          "block": ("<!--", "-->")},
    ".json":   {"name": "JSON"},
    ".md":     {"name": "Markdown"},
    ".dart":   {"name": "Dart",         "line": "//",   "block": ("/*", "*/")},
    ".vue":    {"name": "Vue",          "line": "//",   "block": ("/*", "*/")},
    ".svelte": {"name": "Svelte",       "line": "//",   "block": ("/*", "*/")},
    ".ex":     {"name": "Elixir",       "line": "#"},
    ".exs":    {"name": "Elixir Script","line": "#"},
    ".erl":    {"name": "Erlang",       "line": "%"},
    ".hs":     {"name": "Haskell",      "line": "--",   "block": ("{-", "-}")},
    ".ml":     {"name": "OCaml",        "block": ("(*", "*)")},
    ".clj":    {"name": "Clojure",      "line": ";"},
    ".nim":    {"name": "Nim",          "line": "#"},
    ".zig":    {"name": "Zig",          "line": "//"},
    ".tf":     {"name": "Terraform",    "line": "#",    "block": ("/*", "*/")},
    ".proto":  {"name": "Protobuf",     "line": "//",   "block": ("/*", "*/")},
    ".pine":   {"name": "Pine Script",  "line": "//"},
    ".sol":    {"name": "Solidity",     "line": "//",   "block": ("/*", "*/")},
}

DEFAULT_IGNORE = {
    "node_modules", ".git", "__pycache__", ".venv", "venv", "env", ".env",
    ".tox", ".mypy_cache", ".pytest_cache", "dist", "build", ".next",
    ".nuxt", "target", "bin", "obj", ".idea", ".vscode", ".DS_Store",
    "vendor", "coverage", ".cache", ".turbo", "*.min.js", "*.min.css",
}

def load_gitignore(directory):
    """Load .gitignore patterns from directory."""
    patterns = set()
    gitignore = os.path.join(directory, ".gitignore")
    if os.path.isfile(gitignore):
        with open(gitignore, "r", errors="replace") as f:
            for line in f:
                line = line.strip()
                if line and not line.startswith("#"):
                    patterns.add(line.rstrip("/"))
    return patterns

def should_ignore(path, name, ignore_patterns):
    """Check if a path should be ignored."""
    if name in DEFAULT_IGNORE:
        return True
    for pat in ignore_patterns:
        if fnmatch.fnmatch(name, pat):
            return True
        if fnmatch.fnmatch(path, pat):
            return True
    return False

def count_lines(filepath, lang_info):
    """Count code, comment, and blank lines in a file."""
    code = comments = blanks = 0
    in_block = False
    block_end = None
    line_comment = lang_info.get("line")
    block_start_end = lang_info.get("block")
    alt_block = lang_info.get("alt_block")

    try:
        with open(filepath, "r", errors="replace") as f:
            for raw_line in f:
                line = raw_line.strip()
                if not line:
                    blanks += 1
                    continue
                if in_block:
                    comments += 1
                    if block_end and block_end in line:
                        in_block = False
                    continue
                if block_start_end and line.startswith(block_start_end[0]):
                    comments += 1
                    if block_start_end[1] not in line[len(block_start_end[0]):]:
                        in_block = True
                        block_end = block_start_end[1]
                    continue
                if alt_block and line.startswith(alt_block[0]):
                    comments += 1
                    if alt_block[1] not in line[len(alt_block[0]):]:
                        in_block = True
                        block_end = alt_block[1]
                    continue
                if line_comment and line.startswith(line_comment):
                    comments += 1
                    continue
                code += 1
    except (OSError, UnicodeDecodeError):
        pass
    return code, comments, blanks

def scan_directory(directory, ignore_patterns=None):
    """Scan directory and count lines by language."""
    if ignore_patterns is None:
        ignore_patterns = set()
    gitignore = load_gitignore(directory)
    all_patterns = ignore_patterns | gitignore

    stats = {}  # lang_name -> {files, code, comments, blanks}
    file_details = []

    for root, dirs, files in os.walk(directory):
        dirs[:] = [d for d in dirs if not should_ignore(os.path.join(root, d), d, all_patterns)]
        for fname in files:
            if should_ignore(os.path.join(root, fname), fname, all_patterns):
                continue
            ext = os.path.splitext(fname)[1].lower()
            if ext not in LANGUAGES:
                continue
            lang = LANGUAGES[ext]
            filepath = os.path.join(root, fname)
            code, comments, blanks = count_lines(filepath, lang)
            total = code + comments + blanks
            if total == 0:
                continue
            name = lang["name"]
            if name not in stats:
                stats[name] = {"files": 0, "code": 0, "comments": 0, "blanks": 0}
            stats[name]["files"] += 1
            stats[name]["code"] += code
            stats[name]["comments"] += comments
            stats[name]["blanks"] += blanks
            file_details.append({
                "path": os.path.relpath(filepath, directory),
                "language": name,
                "code": code, "comments": comments, "blanks": blanks,
                "total": total,
            })
    return stats, file_details

def format_table(stats, sort_by="code", title=None):
    """Format stats as a table."""
    rows = []
    for lang, s in stats.items():
        total = s["code"] + s["comments"] + s["blanks"]
        rows.append((lang, s["files"], s["code"], s["comments"], s["blanks"], total))
    sort_idx = {"language": 0, "files": 1, "code": 2, "comments": 3, "blanks": 4, "total": 5}
    idx = sort_idx.get(sort_by, 2)
    rows.sort(key=lambda r: r[idx], reverse=(idx != 0))

    hdr = f"{'Language':<20} {'Files':>7} {'Code':>10} {'Comments':>10} {'Blanks':>10} {'Total':>10}"
    sep = "-" * len(hdr)
    lines = []
    if title:
        lines.append(f"\n  {title}")
    lines.append(sep)
    lines.append(hdr)
    lines.append(sep)
    t_files = t_code = t_comments = t_blanks = t_total = 0
    for lang, files, code, comments, blanks, total in rows:
        lines.append(f"{lang:<20} {files:>7} {code:>10,} {comments:>10,} {blanks:>10,} {total:>10,}")
        t_files += files
        t_code += code
        t_comments += comments
        t_blanks += blanks
        t_total += total
    lines.append(sep)
    lines.append(f"{'Total':<20} {t_files:>7} {t_code:>10,} {t_comments:>10,} {t_blanks:>10,} {t_total:>10,}")
    lines.append(sep)
    return "\n".join(lines)

def format_csv(stats):
    """Format stats as CSV."""
    lines = ["Language,Files,Code,Comments,Blanks,Total"]
    for lang, s in sorted(stats.items()):
        total = s["code"] + s["comments"] + s["blanks"]
        lines.append(f"{lang},{s['files']},{s['code']},{s['comments']},{s['blanks']},{total}")
    return "\n".join(lines)

def format_markdown(stats, title="Code Statistics"):
    """Format stats as a Markdown table."""
    lines = [f"## {title}", "", "| Language | Files | Code | Comments | Blanks | Total |",
             "|----------|------:|-----:|---------:|-------:|------:|"]
    t_files = t_code = t_comments = t_blanks = t_total = 0
    for lang, s in sorted(stats.items(), key=lambda x: x[1]["code"], reverse=True):
        total = s["code"] + s["comments"] + s["blanks"]
        lines.append(f"| {lang} | {s['files']} | {s['code']:,} | {s['comments']:,} | {s['blanks']:,} | {total:,} |")
        t_files += s["files"]
        t_code += s["code"]
        t_comments += s["comments"]
        t_blanks += s["blanks"]
        t_total += total
    lines.append(f"| **Total** | **{t_files}** | **{t_code:,}** | **{t_comments:,}** | **{t_blanks:,}** | **{t_total:,}** |")
    return "\n".join(lines)

def tree_view(file_details, max_depth=3):
    """Generate a directory tree view with line counts."""
    tree = {}
    for f in file_details:
        parts = Path(f["path"]).parts[:max_depth]
        node = tree
        for p in parts:
            if p not in node:
                node[p] = {"_code": 0, "_files": 0, "_children": {}}
            node[p]["_code"] += f["code"]
            node[p]["_files"] += 1
            node = node[p]["_children"]

    lines = []
    def render(node, prefix="", is_last=True):
        items = sorted(node.items())
        for i, (name, data) in enumerate(items):
            last = i == len(items) - 1
            connector = "--- " if last else "|-- "
            code = data["_code"]
            files = data["_files"]
            lines.append(f"{prefix}{connector}{name}/ ({files} files, {code:,} lines)")
            ext = "    " if last else "|   "
            render(data["_children"], prefix + ext, last)
    render(tree)
    return "\n".join(lines)

def compare_dirs(dir1, dir2, sort_by="code"):
    """Compare line counts between two directories."""
    stats1, _ = scan_directory(dir1)
    stats2, _ = scan_directory(dir2)
    all_langs = sorted(set(list(stats1.keys()) + list(stats2.keys())))

    hdr = f"{'Language':<18} {'Dir1 Code':>10} {'Dir2 Code':>10} {'Diff':>10} {'Change':>8}"
    sep = "-" * len(hdr)
    lines = [sep, hdr, sep]
    for lang in all_langs:
        c1 = stats1.get(lang, {}).get("code", 0)
        c2 = stats2.get(lang, {}).get("code", 0)
        diff = c2 - c1
        pct = f"{(diff/c1*100):+.1f}%" if c1 > 0 else ("new" if c2 > 0 else "-")
        lines.append(f"{lang:<18} {c1:>10,} {c2:>10,} {diff:>+10,} {pct:>8}")
    lines.append(sep)
    return "\n".join(lines)

def main():
    parser = argparse.ArgumentParser(
        description="CodeCount - Count lines of code, comments, and blanks by language.",
        formatter_class=argparse.RawDescriptionHelpFormatter,
        epilog="Examples:\n"
               "  %(prog)s .\n"
               "  %(prog)s src/ --json\n"
               "  %(prog)s . --sort files --tree\n"
               "  %(prog)s --compare dir1 dir2\n"
               "  %(prog)s . --markdown > stats.md\n"
    )
    parser.add_argument("directory", nargs="?", default=".", help="Directory to scan (default: .)")
    parser.add_argument("-j", "--json", action="store_true", help="Output as JSON")
    parser.add_argument("-c", "--csv", action="store_true", help="Output as CSV")
    parser.add_argument("-m", "--markdown", action="store_true", help="Output as Markdown table")
    parser.add_argument("-s", "--sort", default="code", choices=["language", "files", "code", "comments", "blanks", "total"], help="Sort by column (default: code)")
    parser.add_argument("-t", "--tree", action="store_true", help="Show directory tree with line counts")
    parser.add_argument("--tree-depth", type=int, default=3, help="Max tree depth (default: 3)")
    parser.add_argument("--compare", nargs=2, metavar=("DIR1", "DIR2"), help="Compare two directories")
    parser.add_argument("-i", "--ignore", nargs="*", default=[], help="Additional patterns to ignore")
    args = parser.parse_args()

    if args.compare:
        print(compare_dirs(args.compare[0], args.compare[1], args.sort))
        return

    directory = os.path.abspath(args.directory)
    if not os.path.isdir(directory):
        print(f"Error: '{directory}' is not a directory.", file=sys.stderr)
        sys.exit(1)

    ignore = set(args.ignore)
    stats, file_details = scan_directory(directory, ignore)

    if not stats:
        print("No source files found.")
        sys.exit(0)

    if args.json:
        output = {"directory": directory, "languages": stats, "files": file_details}
        print(json.dumps(output, indent=2))
    elif args.csv:
        print(format_csv(stats))
    elif args.markdown:
        print(format_markdown(stats, title=os.path.basename(directory)))
    else:
        print(format_table(stats, args.sort, title=os.path.basename(directory)))

    if args.tree:
        print(f"\nDirectory Tree:")
        print(tree_view(file_details, args.tree_depth))

if __name__ == "__main__":
    main()