calculate the complexity of your python files

complexity.py

import os
import ast
from argparse import ArgumentParser, Namespace

COLORS = {
    "red": "\033[91m",
    "green": "\033[92m",
    "yellow": "\033[93m",
    "blue": "\033[94m",
    "magenta": "\033[95m",
    "cyan": "\033[96m",
    "white": "\033[97m",
    "reset": "\033[0m",
}


def compute_complexity(node, depth=0) -> int:
    """
    Compute an approximate cognitive complexity of a python function node.
    - Base complexity: 1
    - Only control-flow constructs add complexity
    - Nested constructs add extra weight proportional to depth
    """
    complexity = 1  # base for function itself

    for child in ast.iter_child_nodes(node):
        if isinstance(child, (ast.If, ast.For, ast.While, ast.With, ast.Match)):
            # decision point → add base 1 + nested penalty
            complexity += 1 + depth
            # recurse for inner control flow
            complexity += compute_complexity(child, depth + 1)
        elif isinstance(child, ast.Try):
            # base try + each except handler
            complexity += 1 + len(child.handlers) + depth
            for sub in ast.iter_child_nodes(child):
                complexity += compute_complexity(sub, depth + 1)
        elif isinstance(child, ast.BoolOp):
            # each extra and/or operation value adds small complexity
            complexity += len(child.values) - 1
        elif isinstance(
            child, (ast.ListComp, ast.DictComp, ast.SetComp, ast.GeneratorExp)
        ):
            complexity += len(child.generators)
        elif isinstance(child, (ast.FunctionDef, ast.ClassDef)):
            # skip nested functions/classes
            continue
        # else: ignore trivial nodes

    return complexity


def analyze_file(path: str, threshold: int = 20):
    with open(path, "r") as f:
        tree = ast.parse(f.read(), filename=path)

    node_types = (ast.FunctionDef, ast.AsyncFunctionDef, ast.ClassDef, ast.Lambda)

    for node in tree.body:
        if isinstance(node, node_types):
            score = compute_complexity(node)
            if score > threshold:
                msg = f"{COLORS['yellow']}[!]{COLORS['reset']} {COLORS['green']}{path}{COLORS['reset']}"
                msg += f":{COLORS['magenta']}{node.name}(){COLORS['reset']} has complexity {COLORS['red']}{score}{COLORS['reset']}"
                print(msg)


def parse_args() -> Namespace:
    parser = ArgumentParser(description="Analyze Python files for cognitive complexity")
    parser.add_argument(
        "-t",
        "--threshold",
        type=int,
        default=20,
        help="Complexity threshold to report (default: 20)",
    )
    parser.add_argument(
        "-r",
        "--recursive",
        default=False,
        action="store_true",
    )
    parser.add_argument(
        "files",
        nargs="*",
        help="Files to analyze (default: all .py files in current directory)",
    )
    return parser.parse_args()


def recursive_file_search(directory: str) -> list[str]:
    retval = []
    for root, _, files in os.walk(directory):
        for file in files:
            if file.endswith(".py"):
                retval.append(os.path.join(root, file))
    return retval


def get_all_python_files(args: Namespace) -> list[str]:
    files = []

    # Default to all files in the current directory if none specified
    if not args.files:
        args.files = [f for f in os.listdir(os.getcwd())]

    # collect python files from the specified paths
    for filename in args.files:
        if os.path.isdir(filename) and args.recursive:
            files.extend(recursive_file_search(filename))
        if os.path.isfile(filename) and filename.endswith(".py"):
            files.append(filename)

    return files


def Main():
    args = parse_args()
    files = get_all_python_files(args)

    for filename in files:
        try:
            analyze_file(filename, threshold=args.threshold)
        except SyntaxError as e:
            print(f"❌ {filename}: failed to parse ({e})")


if __name__ == "__main__":
    Main()