Skip to content

Python:ast

Get docstring of class attribute

클래스의 Attribute 에 추가한 docstring 주석이 잘 획득되지 않는다.

예를 들어 다음과 같은 코드에서:

@dataclass
class StitcherProps:
    stitcher_mode_index: int = 0
    """Scenario for stitcher operation."""

StitcherProps.stitcher_mode_index의 주석은 inspect.getdoc, __doc__ 로 획득되지 않는다.

이 때 사용하는 방법은 ast 모듈을 사용하는 방법이 있다:

# -*- coding: utf-8 -*-

from ast import AnnAssign, Expr, parse
from inspect import cleandoc, getsource
from typing import Optional


def get_attribute_docstring(cls: type, key: str) -> Optional[str]:
    src = getsource(cls)
    module = parse(src)
    module_body = module.body

    assert len(module_body) == 1
    module_body0 = module_body[0]

    # noinspection PyUnresolvedReferences
    assert module_body0.name == cls.__name__

    # noinspection PyUnresolvedReferences
    stmts = module_body0.body
    assert isinstance(stmts, list)

    for i, stmt in enumerate(stmts):
        if not isinstance(stmt, AnnAssign):
            continue
        if stmt.target.id != key:
            continue

        try:
            expr = stmts[i + 1]
        except IndexError:
            return None

        if not isinstance(expr, Expr):
            return None

        # noinspection PyUnresolvedReferences
        doc_value = expr.value.value
        if not isinstance(doc_value, str):
            return None

        return cleandoc(doc_value)

    return None