Auto generate document for a Python Enum class

enum_auto_doc.md

A decorator to document Python Enum value

Why?

Python does not support add program-readable document string to a Enum class value currently, see:

• https://stackoverflow.com/questions/52062831
• https://stackoverflow.com/questions/19330460

How?

There is a approach is to parse AST of the wanted Enum class, and manually add them to document string.

Example

Code

@enum_auto_doc
class SearchType(IntEnum):
    """搜索内容类型"""

    SONG = 1
    """单曲"""
    ALBUM = 10
    """专辑"""
    ARTIST = 100
    """歌手"""
    PLAYLIST = 1000
    """歌单"""
    USER = 1002
    """用户"""
    MV = 1004
    """MV"""
    LYRICS = 1006
    """歌词"""
    DJ = 1009
    """主播电台"""
    VIDEO = 1014
    """视频"""

Result of SearchType.__doc__

搜索内容类型
- `SONG` = *`1`* : 单曲
- `ALBUM` = *`10`* : 专辑
- `ARTIST` = *`100`* : 歌手
- `PLAYLIST` = *`1000`* : 歌单
- `USER` = *`1002`* : 用户
- `MV` = *`1004`* : MV
- `LYRICS` = *`1006`* : 歌词
- `DJ` = *`1009`* : 主播电台
- `VIDEO` = *`1014`* : 视频

TODOs

• Add support for Python version 3.6+ (3.8+ currently)
• Enum class with indentaion support
• Add a AutoEnum metaclass to support auto Enum document by inherit

enum_auto_doc.py

import ast
import inspect
from enum import Enum
from typing import Dict, Type, 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