folding-studio/docs/generate_cli_docs.py

import re
import types
from typing import Union, get_args, get_origin, get_type_hints

import typer
from folding_studio.cli import app
from typer.models import TyperInfo
from typer.utils import get_params_from_function


def extract_base_type(annotation):
    """Extract annotation type.

    Handles List[X], Optional[X] and Union[X, Y] cases.
    """
    origin = get_origin(annotation)

    if origin is list:
        inner_type = get_args(annotation)[0]
        return (
            f"List[{inner_type.__name__}]"
            if hasattr(inner_type, "__name__")
            else "List[Unknown]"
        )

    if origin is Union or isinstance(origin, types.UnionType):
        inner_types = [t.__name__ for t in get_args(annotation) if t is not type(None)]
        return " | ".join(inner_types) if inner_types else "Unknown"

    if isinstance(annotation, types.UnionType):
        inner_types = [t.__name__ for t in get_args(annotation) if t is not type(None)]
        return " | ".join(inner_types) if inner_types else "Unknown"

    return annotation.__name__ if hasattr(annotation, "__name__") else "Unknown"


def update_params_with_default(param_info: dict[str, str]) -> dict[str, str]:
    """
    Update the 'default' value in the params dictionary based on its description.

    Args:
        param_info (dict): A dictionary describing a parameter.

    Returns:
        dict: The updated param_info dictionary.
    """
    if param_info.get("default") != "No default":
        return param_info
    description = param_info.get("description", "")
    # Regular expression to find a sentence starting with 'Default to ' and ending with a period
    match = re.search(r"Default to '([^']+)'", description)
    if not match:
        return param_info
    # Extract the default value
    default_value = match.group(1)
    # Update the 'description' by removing the 'Default to ' phrase
    param_info["description"] = re.sub(
        r"Default to '[^']+'\.\s*", "", description
    ).strip()
    # Ensure the description ends with a single period
    if not param_info["description"].endswith("."):
        param_info["description"] += "."
    # Update the 'default' key with the extracted value
    param_info["default"] = default_value
    param_info["default"] = param_info["default"].replace("<", "&lt;")
    param_info["default"] = param_info["default"].replace(">", "&gt;")
    return param_info


def extract_command_info(command: typer.models.CommandInfo):
    func = command.callback
    parameters = get_params_from_function(func)

    hints = get_type_hints(func)
    docstring = func.__doc__ or "No docstring provided."
    docstring = docstring.replace("\n", " ").replace("\n\n", " ")
    docstring = re.sub(r"(https?://\S+)", r"<\1>", docstring)
    command_info = {
        "docstring": docstring,
        "name": command.name if command.name else func.__name__,
        "params": [],
        "options": [],
    }

    for name, param in parameters.items():
        raw_type = hints.get(name, None)
        base_type = extract_base_type(raw_type)
        description = str(param.default.help)
        description = description.replace("\n", " ").replace("\n\n", " ")
        param_info = {
            "name": name,
            "type": base_type,
            "description": description,
        }

        if isinstance(param.default, typer.models.OptionInfo):
            name = name.replace("_", "-")
            if base_type == "bool":
                name = name + f" / --no-{name}"
            name = f"--{name}"
            param_info["name"] = name
            default_value = param.default.default
            if param_info["type"] == "bool":
                values = param_info["name"].split(" / ")
                default_value = values[0] if param.default.default else values[1]
            if default_value is Ellipsis:
                default_value = "No default"
            param_info["default"] = default_value

            # if No default, check if default value is in description
            param_info = update_params_with_default(param_info)

            command_info["options"].append(param_info)
        else:
            command_info["params"].append(param_info)

    return command_info


def generate_markdown_level_docs(
    f, group: TyperInfo, cli_name: str, level=1, base_name=None
):
    base_level = "#" * level
    f.write(f"{base_level} `{base_name + ' ' if base_name else ''}{group.name}`\n")

    for subcommand in group.typer_instance.registered_commands:
        subcommand_info = extract_command_info(subcommand)
        subcommand_name = (
            subcommand.name
            if subcommand.name is not None
            else subcommand.callback.__name__
        )
        command_name = (
            f"{base_name + ' ' if base_name else ''}{group.name} {subcommand_name}"
        )
        f.write(f"{base_level}# `{command_name}`\n\n")

        f.write(f"{subcommand_info['docstring']}\n\n")

        usage = "**Usage**:\n\n"
        usage += "```console\n"
        usage += f"{cli_name} {command_name}{' [OPTIONS]' if subcommand_info['options'] else ''}"
        if subcommand_info["params"]:
            usage += f" {' '.join(param['name'].upper() for param in subcommand_info['params'])}"
        usage += "\n```\n\n"
        f.write(usage)

        if subcommand_info["params"]:
            # Arguments
            f.write("**Arguments**:\n\n")
            f.write("| ARGUMENT | DESCRIPTION | VALUE TYPE |\n")
            f.write("| -------- | ----------- | ----------- |\n")
            for param in subcommand_info["params"]:
                param["description"] = (
                    param["description"]
                    if param.get("description")
                    else "No description"
                )
                param["type"] = param["type"] if param["type"] else "No type"

                f.write(
                    f"| {param['name'].upper()} | {param['description']} | {param['type']} |\n"
                )

            f.write("\n")

        # Options
        if subcommand_info["options"]:
            f.write("**Options**:\n\n")
            f.write("| OPTIONS | DESCRIPTION | VALUE TYPE | DEFAULT VALUE |\n")
            f.write("| ------- | ----------- | ---------- | ------------- |\n")
            for param in subcommand_info["options"]:
                param["description"] = (
                    param["description"]
                    if param.get("description")
                    else "No description"
                )
                param["type"] = param["type"] if param["type"] else "No type"
                param["default"] = (
                    param["default"] if param["default"] is not None else "No default"
                )
                # if No default, check if default value is in description
                param = update_params_with_default(param)
                f.write(
                    f"| {param['name']} | {param['description']} | {param['type']} | {param['default']} |\n"
                )

            f.write("\n")
    for subgroup in group.typer_instance.registered_groups:
        generate_markdown_level_docs(f, subgroup, cli_name, level + 1, group.name)


def generate_markdown_docs() -> None:
    """
    Generate markdown documentation for all registered commands and subcommands in the application.
    The documentation will include descriptions, arguments, and options.

    The generated markdown is saved in the 'docs/reference/cli.md' file.
    """
    # Iterate through each group in app.registered_groups
    with open("docs/reference/cli.md", "w") as f:
        for group in app.registered_groups:
            if group.name == "key":
                continue
            generate_markdown_level_docs(f, group, "folding", 2)


if __name__ == "__main__":
    generate_markdown_docs()