catch missing or incomplete docstrings and typing in python scripts. enforce standards on LLM-generated code.

doc_check.py

#!/usr/bin/env python3
"""
CLI tool to check Python files for missing docstrings, type hints, and lazy docstrings.

A 'lazy' docstring is one that only contains a summary without documenting
arguments or return values when the function has 2+ parameters or returns something.

Usage:
    python doc_check.py src/ --check-lazy-docstrings

    # CLI help text
    python doc_check.py --help

Dependencies:
    None. stdlib only

Note: docstrings in this script are indeed overkill, 'practice what you preach' etc
"""

import argparse
import ast
import re
import sys
from dataclasses import dataclass, field
from enum import Enum
from pathlib import Path
from typing import Iterator

# ---------------------------------------------------------------------------
# Terminal color helpers + fallback
# ---------------------------------------------------------------------------

_USE_COLOR = sys.stdout.isatty()


def _c(code: str, text: str) -> str:
    """Wrap text in an ANSI escape sequence, or return plain text if color is disabled.

    Args:
        code: ANSI escape code (e.g. '91' for bright red).
        text: The string to format.

    Returns:
        ANSI-formatted string, or the original text if color is disabled.
    """
    if not _USE_COLOR:
        return text
    return f"\033[{code}m{text}\033[0m"


def red(t: str) -> str:
    """Return text formatted in bright red.

    Args:
        t: The string to format.

    Returns:
        ANSI bright-red formatted string.
    """
    return _c("91", t)


def yellow(t: str) -> str:
    """Return text formatted in bright yellow.

    Args:
        t: The string to format.

    Returns:
        ANSI bright-yellow formatted string.
    """
    return _c("93", t)


def green(t: str) -> str:
    """Return text formatted in bright green.

    Args:
        t: The string to format.

    Returns:
        ANSI bright-green formatted string.
    """
    return _c("92", t)


def cyan(t: str) -> str:
    """Return text formatted in bright cyan.

    Args:
        t: The string to format.

    Returns:
        ANSI bright-cyan formatted string.
    """
    return _c("96", t)


def bold(t: str) -> str:
    """Return text formatted in bold.

    Args:
        t: The string to format.

    Returns:
        ANSI bold formatted string.
    """
    return _c("1", t)


def dim(t: str) -> str:
    """Return text formatted in dim/faint.

    Args:
        t: The string to format.

    Returns:
        ANSI dim formatted string.
    """
    return _c("2", t)


# Map issue types to colors
_ISSUE_COLORS = {
    "docstring": yellow,
    "lazy_docstring": yellow,
    "return_type": red,
    "param_type": red,
    "syntax": red,
    "error": red,
}


def _pretty_path(path: Path) -> str:
    """Return a human-readable absolute path, using ~/... when under home.

    Args:
        path: Resolved absolute path to display.

    Returns:
        Path string with home directory replaced by '~' where applicable.
    """
    try:
        return "~/" + str(path.relative_to(Path.home()))
    except ValueError:
        return str(path)


def _display_path(target: Path, filepath: Path) -> str:
    """Return a clean display path for a file relative to the scan root.

    When the file is directly under the scan root, shows only the filename
    or short relative path. Otherwise falls back to a home-relative or
    absolute path — never a '../../..' chain.

    Args:
        target: The scan root (file or directory) used as the relative base.
        filepath: The file whose display path is being computed.

    Returns:
        A concise string path suitable for terminal output.
    """
    root = target if target.is_dir() else target.parent
    try:
        return str(filepath.relative_to(root))
    except ValueError:
        return _pretty_path(filepath)


# ---------------------------------------------------------------------------
# Data models
# ---------------------------------------------------------------------------


class CheckMode(Enum):
    """Available checking modes."""

    DOCSTRING = "docstring"
    TYPING = "typing"
    BOTH = "both"


@dataclass
class Issue:
    """Represents a single linting issue."""

    filepath: Path
    name: str
    line: int
    issue_type: str
    detail: str


@dataclass
class CheckResult:
    """Aggregated results for a file."""

    filepath: Path
    issues: list[Issue] = field(default_factory=list)


# ---------------------------------------------------------------------------
# Docstring parsing
# ---------------------------------------------------------------------------


@dataclass
class DocstringInfo:
    """Parsed information from a docstring."""

    raw: str
    summary: str = ""
    body: str = ""
    documented_params: set[str] = field(default_factory=set)
    has_return_doc: bool = False

    @classmethod
    def parse(cls, docstring: str | None) -> "DocstringInfo | None":
        """Parse a docstring and extract structured information.

        Style-agnostic: detects Google, NumPy, Sphinx, and freeform styles
        by checking if param names actually appear with descriptions.

        Args:
            docstring: The raw docstring text to parse.

        Returns:
            DocstringInfo object with parsed details, or None if docstring is None.
        """
        if docstring is None:
            return None

        lines = docstring.strip().split("\n")

        summary_lines: list[str] = []
        body_start = 0
        for i, line in enumerate(lines):
            if line.strip() == "":
                body_start = i + 1
                break
            summary_lines.append(line)
        else:
            body_start = len(lines)

        summary = " ".join(summary_lines)
        body = "\n".join(lines[body_start:])

        info = cls(raw=docstring, summary=summary, body=body)
        info.documented_params = cls._extract_documented_params(docstring, body)
        info.has_return_doc = cls._has_return_documentation(body)
        return info

    @staticmethod
    def _extract_documented_params(full_doc: str, body: str) -> set[str]:
        """Extract parameter names that appear to be documented.

        Args:
            full_doc: Complete docstring text.
            body: Docstring body (after summary paragraph).

        Returns:
            Set of lowercase parameter names that are documented.
        """
        documented: set[str] = set()

        # Sphinx: :param name: or :param type name:
        for match in re.finditer(r":param\s+([^:]+):", full_doc):
            parts = match.group(1).strip().split()
            if parts:
                documented.add(parts[-1].lower())

        # Google: indented `name:` or `name (type):` with description
        for match in re.finditer(
            r"^\s{2,}(\*{0,2}\w+)\s*(?:\([^)]*\))?\s*:", body, re.MULTILINE
        ):
            name = match.group(1).lstrip("*")
            documented.add(name.lower())

        # NumPy: `name : type`
        for match in re.finditer(r"^\s*(\w+)\s+:\s+\S+", body, re.MULTILINE):
            documented.add(match.group(1).lower())

        return documented

    @staticmethod
    def _has_return_documentation(body: str) -> bool:
        """Check if return value is documented in the body.

        Args:
            body: Docstring body (after summary paragraph).

        Returns:
            True if any return/yield documentation pattern is found.
        """
        patterns = [
            r":returns?\s*[^:]*:",
            r":rtype\s*:",
            r"^\s*returns?\s*[:\-]",
            r"^\s*returns?\s*\n\s*-{3,}",
            r"^\s*yields?\s*[:\-]",
            r"^\s*yields?\s*\n\s*-{3,}",
        ]
        for pattern in patterns:
            if re.search(pattern, body, re.MULTILINE | re.IGNORECASE):
                return True
        return False

    def is_param_documented(self, param_name: str) -> bool:
        """Check if a specific parameter is documented.

        Args:
            param_name: Name of the parameter to check.

        Returns:
            True if the parameter appears in documented_params.
        """
        return param_name.lstrip("*").lower() in self.documented_params


# ---------------------------------------------------------------------------
# AST checker
# ---------------------------------------------------------------------------


class DocstringTypingChecker(ast.NodeVisitor):
    """AST visitor that checks for missing docstrings, type hints, and lazy docstrings."""

    def __init__(
        self, filepath: Path, mode: CheckMode, check_lazy: bool = False
    ) -> None:
        """Initialize checker with filepath and mode.

        Args:
            filepath: Path to the file being checked.
            mode: The checking mode (docstring, typing, or both).
            check_lazy: Whether to check for lazy/insufficient docstrings.
        """
        self.filepath = filepath
        self.mode = mode
        self.check_lazy = check_lazy
        self.issues: list[Issue] = []
        self._class_stack: list[str] = []

    def _qualified_name(self, name: str) -> str:
        """Return fully qualified name including parent class if any.

        Args:
            name: The base name of the function or class.

        Returns:
            Qualified name with parent class prefix if inside a class.
        """
        if self._class_stack:
            return f"{self._class_stack[-1]}.{name}"
        return name

    def _should_check_docstrings(self) -> bool:
        """Return True if docstring checks are enabled for the current mode.

        Returns:
            True when mode is DOCSTRING or BOTH.
        """
        return self.mode in (CheckMode.DOCSTRING, CheckMode.BOTH)

    def _should_check_typing(self) -> bool:
        """Return True if type hint checks are enabled for the current mode.

        Returns:
            True when mode is TYPING or BOTH.
        """
        return self.mode in (CheckMode.TYPING, CheckMode.BOTH)

    def _add_issue(self, name: str, line: int, issue_type: str, detail: str) -> None:
        """Record an issue found during checking.

        Args:
            name: Qualified name of the entity with the issue.
            line: Line number where the issue occurs.
            issue_type: Category of the issue (e.g. docstring, param_type).
            detail: Human-readable description of the problem.
        """
        self.issues.append(
            Issue(
                filepath=self.filepath,
                name=name,
                line=line,
                issue_type=issue_type,
                detail=detail,
            )
        )

    def _get_real_params(
        self, node: ast.FunctionDef | ast.AsyncFunctionDef
    ) -> list[str]:
        """Extract parameter names excluding self/cls.

        Args:
            node: The function definition AST node.

        Returns:
            List of parameter names, with vararg/kwarg prefixed by */**.
        """
        args = node.args
        all_args = args.posonlyargs + args.args + args.kwonlyargs
        params = [arg.arg for arg in all_args if arg.arg not in ("self", "cls")]
        if args.vararg:
            params.append(f"*{args.vararg.arg}")
        if args.kwarg:
            params.append(f"**{args.kwarg.arg}")
        return params

    def _function_has_return(
        self, node: ast.FunctionDef | ast.AsyncFunctionDef
    ) -> bool:
        """Check if a function has any non-None return or yield statement.

        Does not descend into nested function definitions.

        Args:
            node: The function definition AST node.

        Returns:
            True if the function contains a meaningful return or yield.
        """
        to_visit = list(ast.iter_child_nodes(node))
        while to_visit:
            child = to_visit.pop()
            if isinstance(child, (ast.FunctionDef, ast.AsyncFunctionDef)):
                continue
            if isinstance(child, ast.Return) and child.value is not None:
                if not (
                    isinstance(child.value, ast.Constant) and child.value.value is None
                ):
                    return True
            if isinstance(child, (ast.Yield, ast.YieldFrom)):
                return True
            to_visit.extend(ast.iter_child_nodes(child))
        return False

    def _check_lazy_docstring(
        self, node: ast.FunctionDef | ast.AsyncFunctionDef, qualified_name: str
    ) -> None:
        """Check if a docstring is lazy/insufficient.

        A docstring is considered lazy if the function has 2+ params but not
        all are documented, or if it returns a value with no Returns section.

        Args:
            node: The function definition AST node.
            qualified_name: The fully qualified name of the function.
        """
        docstring = ast.get_docstring(node)
        if docstring is None:
            return

        doc_info = DocstringInfo.parse(docstring)
        if doc_info is None:
            return

        params = self._get_real_params(node)
        has_return = self._function_has_return(node) and node.name != "__init__"

        issues_found: list[str] = []

        if len(params) >= 2:
            undocumented = [
                p for p in params if not doc_info.is_param_documented(p.lstrip("*"))
            ]
            if undocumented:
                if len(undocumented) == len(params):
                    issues_found.append(f"has {len(params)} params but none documented")
                else:
                    issues_found.append(f"missing docs for: {', '.join(undocumented)}")

        if has_return and not doc_info.has_return_doc:
            issues_found.append("returns value but no Returns section")

        if issues_found:
            self._add_issue(
                qualified_name,
                node.lineno,
                "lazy_docstring",
                "; ".join(issues_found),
            )

    def _check_function(self, node: ast.FunctionDef | ast.AsyncFunctionDef) -> None:
        """Check a function/method for docstring and type hint issues.

        Args:
            node: The function definition AST node to check.
        """
        qualified_name = self._qualified_name(node.name)

        # Skip dunder methods except __init__
        if (
            node.name.startswith("__")
            and node.name.endswith("__")
            and node.name != "__init__"
        ):
            return

        if self._should_check_docstrings() and not ast.get_docstring(node):
            self._add_issue(
                qualified_name, node.lineno, "docstring", "Missing docstring"
            )

        if self.check_lazy:
            self._check_lazy_docstring(node, qualified_name)

        if self._should_check_typing():
            if node.name != "__init__" and node.returns is None:
                self._add_issue(
                    qualified_name,
                    node.lineno,
                    "return_type",
                    "Missing return type hint",
                )

            args = node.args
            all_args = args.posonlyargs + args.args + args.kwonlyargs

            for arg in all_args:
                if arg.arg in ("self", "cls"):
                    continue
                if arg.annotation is None:
                    self._add_issue(
                        qualified_name,
                        getattr(arg, "lineno", node.lineno),
                        "param_type",
                        f"Missing type hint for parameter '{arg.arg}'",
                    )

            if args.vararg and args.vararg.annotation is None:
                self._add_issue(
                    qualified_name,
                    getattr(args.vararg, "lineno", node.lineno),
                    "param_type",
                    f"Missing type hint for *{args.vararg.arg}",
                )
            if args.kwarg and args.kwarg.annotation is None:
                self._add_issue(
                    qualified_name,
                    getattr(args.kwarg, "lineno", node.lineno),
                    "param_type",
                    f"Missing type hint for **{args.kwarg.arg}",
                )

    def visit_FunctionDef(self, node: ast.FunctionDef) -> None:
        """Visit a function definition node.

        Args:
            node: The FunctionDef AST node.
        """
        self._check_function(node)
        self.generic_visit(node)

    def visit_AsyncFunctionDef(self, node: ast.AsyncFunctionDef) -> None:
        """Visit an async function definition node.

        Args:
            node: The AsyncFunctionDef AST node.
        """
        self._check_function(node)
        self.generic_visit(node)

    def visit_ClassDef(self, node: ast.ClassDef) -> None:
        """Visit a class definition node.

        Args:
            node: The ClassDef AST node.
        """
        if self._should_check_docstrings() and not ast.get_docstring(node):
            self._add_issue(
                node.name, node.lineno, "docstring", "Missing class docstring"
            )
        self._class_stack.append(node.name)
        self.generic_visit(node)
        self._class_stack.pop()


# ---------------------------------------------------------------------------
# File checking
# ---------------------------------------------------------------------------


def check_file(
    filepath: Path, mode: CheckMode, check_lazy: bool = False
) -> CheckResult:
    """Parse and check a single Python file.

    Args:
        filepath: Path to the Python file to check.
        mode: The checking mode to use.
        check_lazy: Whether to check for lazy docstrings.

    Returns:
        CheckResult containing all issues found in the file.
    """
    result = CheckResult(filepath=filepath)

    try:
        source = filepath.read_text(encoding="utf-8")
        tree = ast.parse(source, filename=str(filepath))
    except SyntaxError as e:
        result.issues.append(
            Issue(
                filepath, "<module>", e.lineno or 1, "syntax", f"Syntax error: {e.msg}"
            )
        )
        return result
    except Exception as e:
        result.issues.append(
            Issue(filepath, "<module>", 1, "error", f"Parse error: {e}")
        )
        return result

    if mode in (CheckMode.DOCSTRING, CheckMode.BOTH) and not ast.get_docstring(tree):
        result.issues.append(
            Issue(filepath, "<module>", 1, "docstring", "Missing module docstring")
        )

    checker = DocstringTypingChecker(filepath, mode, check_lazy)
    checker.visit(tree)
    result.issues.extend(checker.issues)

    return result


def find_python_files(
    directory: Path, exclude_patterns: list[str] | None = None
) -> Iterator[Path]:
    """Recursively find all Python files in a directory.

    Args:
        directory: Root directory to search.
        exclude_patterns: Directory names to exclude from search.

    Yields:
        Path objects for each Python file found.
    """
    if not directory.is_dir():
        raise ValueError(f"Expected a directory, got: {directory}")

    exclude_patterns = exclude_patterns or []
    skip_dirs = {
        ".venv",
        "venv",
        "__pycache__",
        ".git",
        "node_modules",
        ".tox",
        "build",
        "dist",
    } | set(exclude_patterns)
    skip_files = {"_version.py"}

    for path in directory.rglob("*.py"):
        if any(part in skip_dirs for part in path.parts):
            continue
        if path.name in skip_files:
            continue
        yield path


# ---------------------------------------------------------------------------
# Output formatting
# ---------------------------------------------------------------------------


def format_issues(
    results: list[CheckResult],
    scan_root: Path,
    verbose: bool = False,
) -> str:
    """Format issues for terminal output.

    Args:
        results: List of CheckResult objects to format.
        scan_root: Root used when the scan was initiated (file or directory).
        verbose: Whether to show files without issues.

    Returns:
        Formatted string ready for terminal output.
    """
    lines: list[str] = []
    total_issues = 0
    files_with_issues = 0
    issues_by_type: dict[str, int] = {}

    for result in results:
        dpath = _display_path(scan_root, result.filepath)

        if not result.issues:
            if verbose:
                lines.append(green(f"  OK  {dpath}"))
            continue

        files_with_issues += 1
        total_issues += len(result.issues)

        lines.append(f"\n{bold(dpath)}")
        lines.append(dim("─" * min(len(dpath) + 2, 80)))

        for issue in sorted(result.issues, key=lambda i: i.line):
            issues_by_type[issue.issue_type] = (
                issues_by_type.get(issue.issue_type, 0) + 1
            )
            color_fn = _ISSUE_COLORS.get(issue.issue_type, cyan)
            type_label = color_fn(f"{issue.issue_type:<16}")
            lines.append(
                f"  {dim(f'L{issue.line:4d}')}  {type_label}  "
                f"{cyan(issue.name)}: {issue.detail}"
            )

    # Summary — only print if there's something to say
    lines.append(f"\n{bold('=' * 60)}")
    lines.append(f"Files scanned : {len(results)}")

    if files_with_issues:
        lines.append(red(f"Files w/ issues: {files_with_issues}"))
        lines.append(red(f"Total issues   : {total_issues}"))
        lines.append(f"\n{bold('Breakdown by type:')}")
        for issue_type, count in sorted(issues_by_type.items(), key=lambda x: -x[1]):
            color_fn = _ISSUE_COLORS.get(issue_type, cyan)
            lines.append(f"  {color_fn(f'{issue_type:<16}')}  {count}")
    else:
        lines.append(green("No issues found."))

    return "\n".join(lines)


# ---------------------------------------------------------------------------
# CLI
# ---------------------------------------------------------------------------


def get_parser() -> argparse.ArgumentParser:
    """Build and return the argument parser.

    Returns:
        Configured ArgumentParser instance.
    """
    parser = argparse.ArgumentParser(
        description=(
            "Check Python files for missing docstrings, type hints, and lazy docstrings."
        ),
        formatter_class=argparse.ArgumentDefaultsHelpFormatter,
    )
    parser.add_argument(
        "path",
        type=Path,
        help="File or directory to check",
    )
    parser.add_argument(
        "-m",
        "--mode",
        choices=["docstring", "typing", "both"],
        default="both",
        help="What to check for",
    )
    parser.add_argument(
        "-v",
        "--verbose",
        action="store_true",
        help="Show files without issues",
    )
    parser.add_argument(
        "--exclude",
        nargs="*",
        default=[],
        metavar="DIR",
        help="Directory names to exclude (in addition to defaults)",
    )
    parser.add_argument(
        "--strict",
        action="store_true",
        help="Exit with code 1 if any issues found",
    )
    parser.add_argument(
        "--check-lazy-docstrings",
        "--check-lazy",
        action="store_true",
        dest="check_lazy",
        help="Flag docstrings that lack Args/Returns sections when needed",
    )
    parser.add_argument(
        "--no-color",
        action="store_true",
        dest="no_color",
        help="Disable ANSI color output",
    )
    return parser


def main() -> int:
    """CLI entry point.

    Returns:
        Exit code (0 for success, 1 for issues in strict mode, 2 for errors).
    """
    args = get_parser().parse_args()

    global _USE_COLOR
    if args.no_color:
        _USE_COLOR = False

    target: Path = args.path.resolve()

    if not target.exists():
        print(f"{red('Error:')} path does not exist: {target}", file=sys.stderr)
        return 2

    if target.is_file():
        if target.suffix != ".py":
            print(f"{red('Error:')} not a Python file: {target}", file=sys.stderr)
            return 2
        files = [target]
        scan_root = target
    else:
        files = sorted(find_python_files(target, args.exclude))
        scan_root = target

    if not files:
        print("No Python files found.", file=sys.stderr)
        return 0

    print(
        f"Checking {bold(_pretty_path(target))}  ({len(files)} file{'s' if len(files) != 1 else ''})"
    )

    mode = CheckMode(args.mode)
    results = [check_file(f, mode, args.check_lazy) for f in files]

    print(format_issues(results, scan_root, verbose=args.verbose))

    if args.strict:
        return 1 if any(r.issues for r in results) else 0
    return 0


if __name__ == "__main__":
    sys.exit(main())

example output/use: