hibiapi/utils/decorators/enum.py

import ast
import inspect
from enum import Enum
from typing import TypeVar

_ET = TypeVar("_ET", bound=type[Enum])


def enum_auto_doc(enum: _ET) -> _ET:
    enum_class_ast, *_ = ast.parse(inspect.getsource(enum)).body
    assert isinstance(enum_class_ast, ast.ClassDef)

    enum_value_comments: dict[str, str] = {}
    for index, body in enumerate(body_list := enum_class_ast.body):
        if (
            isinstance(body, ast.Assign)
            and (next_index := index + 1) < len(body_list)
            and isinstance(next_body := body_list[next_index], ast.Expr)
        ):
            target, *_ = body.targets
            assert isinstance(target, ast.Name)
            assert isinstance(next_body.value, ast.Constant)
            assert isinstance(member_doc := next_body.value.value, str)
            enum[target.id].__doc__ = member_doc
            enum_value_comments[target.id] = inspect.cleandoc(member_doc)

    if not enum_value_comments and all(member.name == member.value for member in enum):
        return enum

    members_doc = ""
    for member in enum:
        value_document = "-"
        if member.name != member.value:
            value_document += f" `{member.name}` ="
        value_document += f" *`{member.value}`*"
        if doc := enum_value_comments.get(member.name):
            value_document += f" : {doc}"
        members_doc += value_document + "\n"

    enum.__doc__ = f"{enum.__doc__}\n{members_doc}"
    return enum