1"""Utilities related to attribute docstring extraction.""" 
    2 
    3from __future__ import annotations 
    4 
    5import ast 
    6import inspect 
    7import sys 
    8import textwrap 
    9from typing import Any 
    10 
    11 
    12class DocstringVisitor(ast.NodeVisitor): 
    13    def __init__(self) -> None: 
    14        super().__init__() 
    15 
    16        self.target: str | None = None 
    17        self.attrs: dict[str, str] = {} 
    18        self.previous_node_type: type[ast.AST] | None = None 
    19 
    20    def visit(self, node: ast.AST) -> Any: 
    21        node_result = super().visit(node) 
    22        self.previous_node_type = type(node) 
    23        return node_result 
    24 
    25    def visit_AnnAssign(self, node: ast.AnnAssign) -> Any: 
    26        if isinstance(node.target, ast.Name): 
    27            self.target = node.target.id 
    28 
    29    def visit_Expr(self, node: ast.Expr) -> Any: 
    30        if ( 
    31            isinstance(node.value, ast.Constant) 
    32            and isinstance(node.value.value, str) 
    33            and self.previous_node_type is ast.AnnAssign 
    34        ): 
    35            docstring = inspect.cleandoc(node.value.value) 
    36            if self.target: 
    37                self.attrs[self.target] = docstring 
    38            self.target = None 
    39 
    40 
    41def _dedent_source_lines(source: list[str]) -> str: 
    42    # Required for nested class definitions, e.g. in a function block 
    43    dedent_source = textwrap.dedent(''.join(source)) 
    44    if dedent_source.startswith((' ', '\t')): 
    45        # We are in the case where there's a dedented (usually multiline) string 
    46        # at a lower indentation level than the class itself. We wrap our class 
    47        # in a function as a workaround. 
    48        dedent_source = f'def dedent_workaround():\n{dedent_source}' 
    49    return dedent_source 
    50 
    51 
    52def _extract_source_from_frame(cls: type[Any]) -> list[str] | None: 
    53    frame = inspect.currentframe() 
    54 
    55    while frame: 
    56        if inspect.getmodule(frame) is inspect.getmodule(cls): 
    57            lnum = frame.f_lineno 
    58            try: 
    59                lines, _ = inspect.findsource(frame) 
    60            except OSError:  # pragma: no cover 
    61                # Source can't be retrieved (maybe because running in an interactive terminal), 
    62                # we don't want to error here. 
    63                pass 
    64            else: 
    65                block_lines = inspect.getblock(lines[lnum - 1 :]) 
    66                dedent_source = _dedent_source_lines(block_lines) 
    67                try: 
    68                    block_tree = ast.parse(dedent_source) 
    69                except SyntaxError: 
    70                    pass 
    71                else: 
    72                    stmt = block_tree.body[0] 
    73                    if isinstance(stmt, ast.FunctionDef) and stmt.name == 'dedent_workaround': 
    74                        # `_dedent_source_lines` wrapped the class around the workaround function 
    75                        stmt = stmt.body[0] 
    76                    if isinstance(stmt, ast.ClassDef) and stmt.name == cls.__name__: 
    77                        return block_lines 
    78 
    79        frame = frame.f_back 
    80 
    81 
    82def extract_docstrings_from_cls(cls: type[Any], use_inspect: bool = False) -> dict[str, str]: 
    83    """Map model attributes and their corresponding docstring. 
    84 
    85    Args: 
    86        cls: The class of the Pydantic model to inspect. 
    87        use_inspect: Whether to skip usage of frames to find the object and use 
    88            the `inspect` module instead. 
    89 
    90    Returns: 
    91        A mapping containing attribute names and their corresponding docstring. 
    92    """ 
    93    if use_inspect or sys.version_info >= (3, 13): 
    94        # On Python < 3.13, `inspect.getsourcelines()` might not work as expected 
    95        # if two classes have the same name in the same source file. 
    96        # On Python 3.13+, it will use the new `__firstlineno__` class attribute, 
    97        # making it way more robust. 
    98        try: 
    99            source, _ = inspect.getsourcelines(cls) 
    100        except OSError:  # pragma: no cover 
    101            return {} 
    102    else: 
    103        # TODO remove this implementation when we drop support for Python 3.12: 
    104        source = _extract_source_from_frame(cls) 
    105 
    106    if not source: 
    107        return {} 
    108 
    109    dedent_source = _dedent_source_lines(source) 
    110 
    111    visitor = DocstringVisitor() 
    112    visitor.visit(ast.parse(dedent_source)) 
    113    return visitor.attrs