Coverage for /pythoncovmergedfiles/medio/medio/usr/local/lib/python3.11/site-packages/pydantic/_internal/_docs_extraction.py: 26%

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 # Don't look into nested classes and functions, as their attributes don't belong to the inspected class: 

41 def visit_ClassDef(self, node: ast.ClassDef) -> Any: 

42 pass 

43 

44 def visit_FunctionDef(self, node: ast.FunctionDef) -> Any: 

45 pass 

46 

47 def visit_AsyncFunctionDef(self, node: ast.AsyncFunctionDef) -> Any: 

48 pass 

49 

50 

51def _dedent_source_lines(source: list[str]) -> str: 

52 # Required for nested class definitions, e.g. in a function block 

53 dedent_source = textwrap.dedent(''.join(source)) 

54 if dedent_source.startswith((' ', '\t')): 

55 # We are in the case where there's a dedented (usually multiline) string 

56 # at a lower indentation level than the class itself. We wrap our class 

57 # in a function as a workaround. 

58 dedent_source = f'def dedent_workaround():\n{dedent_source}' 

59 return dedent_source 

60 

61 

62def _extract_source_from_frame(cls: type[Any]) -> list[str] | None: 

63 frame = inspect.currentframe() 

64 

65 while frame: 

66 if inspect.getmodule(frame) is inspect.getmodule(cls): 

67 lnum = frame.f_lineno 

68 try: 

69 lines, _ = inspect.findsource(frame) 

70 except OSError: # pragma: no cover 

71 # Source can't be retrieved (maybe because running in an interactive terminal), 

72 # we don't want to error here. 

73 pass 

74 else: 

75 block_lines = inspect.getblock(lines[lnum - 1 :]) 

76 dedent_source = _dedent_source_lines(block_lines) 

77 try: 

78 block_tree = ast.parse(dedent_source) 

79 except SyntaxError: 

80 pass 

81 else: 

82 stmt = block_tree.body[0] 

83 if isinstance(stmt, ast.FunctionDef) and stmt.name == 'dedent_workaround': 

84 # `_dedent_source_lines` wrapped the class around the workaround function 

85 stmt = stmt.body[0] 

86 if isinstance(stmt, ast.ClassDef) and stmt.name == cls.__name__: 

87 return block_lines 

88 

89 frame = frame.f_back 

90 

91 

92def extract_docstrings_from_cls(cls: type[Any], use_inspect: bool = False) -> dict[str, str]: 

93 """Map model attributes and their corresponding docstring. 

94 

95 Args: 

96 cls: The class of the Pydantic model to inspect. 

97 use_inspect: Whether to skip usage of frames to find the object and use 

98 the `inspect` module instead. 

99 

100 Returns: 

101 A mapping containing attribute names and their corresponding docstring. 

102 """ 

103 if use_inspect or sys.version_info >= (3, 13): 

104 # On Python < 3.13, `inspect.getsourcelines()` might not work as expected 

105 # if two classes have the same name in the same source file. 

106 # On Python 3.13+, it will use the new `__firstlineno__` class attribute, 

107 # making it way more robust. 

108 try: 

109 source, _ = inspect.getsourcelines(cls) 

110 except (OSError, TypeError): # pragma: no cover 

111 return {} 

112 else: 

113 # TODO remove this implementation when we drop support for Python 3.12: 

114 source = _extract_source_from_frame(cls) 

115 

116 if not source: 

117 return {} 

118 

119 dedent_source = _dedent_source_lines(source) 

120 

121 stmt = ast.parse(dedent_source).body[0] 

122 if isinstance(stmt, ast.FunctionDef) and stmt.name == 'dedent_workaround': 

123 # `_dedent_source_lines()` wrapped the class around the workaround function 

124 stmt = stmt.body[0] 

125 

126 visitor = DocstringVisitor() 

127 visitor.generic_visit(stmt) 

128 return visitor.attrs