Skip to content
Search

Python:ast

lastmod 2024-07-04created 2024-07-042 edits

Get docstring of class attribute

Docstring comments added to a class attribute are hard to retrieve.

For example, in the following code:

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

The comment on StitcherProps.stitcher_mode_index cannot be retrieved with inspect.getdoc or __doc__.

In this case, one approach is to use the ast module:

# -*- 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
author Zer0 sitemap