llama-stack-mirror/scripts/check-logger-usage.py

#!/usr/bin/env python3
# Copyright (c) Meta Platforms, Inc. and affiliates.
# All rights reserved.
#
# This source code is licensed under the terms described in the LICENSE file in
# the root directory of this source tree.

"""
Check for improper logger usage in Python files.

This script enforces the use of the custom logger implementation:
    from llama_stack.log import get_logger
    logger = get_logger(name=__name__, category="core")

Instead of the standard logging module:
    import logging

Certain files are excluded where import logging is necessary.
"""

import argparse
import ast
import sys
from pathlib import Path

# Files where import logging is allowed (relative to project root)
ALLOWED_LOGGING_IMPORTS = {
    "llama_stack/log.py",  # The logger implementation itself needs logging
    "log.py",  # When running from within llama_stack directory
}

# Additional patterns for files that may legitimately need logging
ALLOWED_PATTERNS = {
    # Test files might need to directly control logging for testing
    "tests/",
    # Third-party integration code might need direct logging access
    "llama_stack/providers/remote/",
    # Some model implementations might need direct logging control
    "llama_stack/models/",
    # Distribution/build scripts might need direct logging
    "llama_stack/distribution/build.py",
    "llama_stack/distribution/configure.py",
    "llama_stack/distribution/library_client.py",
    "llama_stack/distribution/request_headers.py",
    "llama_stack/distribution/server/server.py",
    "llama_stack/distribution/utils/",
    # Inline providers might need direct logging control
    "llama_stack/providers/inline/",
    # Provider utilities might need direct logging
    "llama_stack/providers/utils/",
}


def is_file_excluded(file_path: str) -> bool:
    """Check if a file is excluded from the logging import check."""
    # Check exact file matches
    if file_path in ALLOWED_LOGGING_IMPORTS:
        return True

    # Check pattern matches
    for pattern in ALLOWED_PATTERNS:
        if file_path.startswith(pattern):
            return True

    return False


def find_logging_imports(file_path: Path) -> list[int]:
    """
    Find line numbers where 'import logging' is used in a Python file.

    Returns:
        List of line numbers where import logging is found.
    """
    try:
        with open(file_path, encoding="utf-8") as f:
            content = f.read()

        try:
            tree = ast.parse(content)
        except SyntaxError:
            # Skip files with syntax errors (might be templates or broken files)
            return []

        logging_imports = []

        for node in ast.walk(tree):
            if isinstance(node, ast.Import):
                for alias in node.names:
                    if alias.name == "logging":
                        logging_imports.append(node.lineno)
            elif isinstance(node, ast.ImportFrom):
                if node.module == "logging":
                    logging_imports.append(node.lineno)

        return logging_imports
    except (OSError, UnicodeDecodeError):
        # Skip files that can't be read
        return []


def check_directory(directory: Path) -> int:
    """
    Check all Python files in a directory for improper logging imports.

    Returns:
        Number of violations found.
    """
    violations = 0

    # Find all Python files
    python_files = directory.rglob("*.py")

    for py_file in python_files:
        # Get relative path for checking against exclusions
        try:
            relative_path = str(py_file.relative_to(directory))
        except ValueError:
            # File is not under the directory (shouldn't happen with rglob)
            continue

        # Skip excluded files
        if is_file_excluded(relative_path):
            continue

        # Check for logging imports
        logging_lines = find_logging_imports(py_file)

        if logging_lines:
            violations += 1
            # Format for GitHub Actions annotations for each logger violation
            for line_num in logging_lines:
                title = "Improper logger usage"
                message = (
                    f"File '{relative_path}' uses improper logger. "
                    "Use the custom logger instead: "
                    "from llama_stack.log import get_logger; "
                    'logger = get_logger(name=__name__, category="core")'
                )
                print(f"::error title={title},file={relative_path},line={line_num}::{message}")

    return violations


def main():
    parser = argparse.ArgumentParser(description="Check for improper logger usage")
    parser.add_argument("directory", nargs="?", default="llama_stack", help="Directory to check (default: llama_stack)")

    args = parser.parse_args()

    # Convert to Path object
    check_dir = Path(args.directory)

    if not check_dir.exists():
        print(f"ERROR: Directory '{check_dir}' does not exist", file=sys.stderr)
        return 1

    if not check_dir.is_dir():
        print(f"ERROR: '{check_dir}' is not a directory", file=sys.stderr)
        return 1

    # Perform the check
    violations = check_directory(check_dir)

    if violations > 0:
        print(f"Found {violations} file(s) with improper logger usage.")
        print()
        print("To fix these issues:")
        print("   1. Replace 'import logging' with 'from llama_stack.log import get_logger'")
        print(
            "   2. Replace logger creation with 'logger = get_logger(name=__name__, category=\"appropriate_category\")'"
        )
        print("   3. Available categories: core, server, router, inference, agents, safety, eval, tools, client")
        print()
        return 1
    else:
        print("All files use proper logger implementation.")
        return 0


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