Spaces:
Runtime error
Runtime error
File size: 1,461 Bytes
0a1b571 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 |
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
|