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