Coverage for pydantic/_internal/_docs_extraction.py: 90.36%

63 statements  

« prev     ^ index     » next       coverage.py v7.5.4, created at 2024-07-03 19:29 +0000

1"""Utilities related to attribute docstring extraction.""" 

2 

3from __future__ import annotations 1abcdefghABCDijklmnopqrEFGHMNOPQRSTUVstuvwxyzIJKL

4 

5import ast 1abcdefghABCDijklmnopqrEFGHMNOPQRSTUVstuvwxyzIJKL

6import inspect 1abcdefghABCDijklmnopqrEFGHMNOPQRSTUVstuvwxyzIJKL

7import textwrap 1abcdefghABCDijklmnopqrEFGHMNOPQRSTUVstuvwxyzIJKL

8from typing import Any 1abcdefghABCDijklmnopqrEFGHMNOPQRSTUVstuvwxyzIJKL

9 

10 

11class DocstringVisitor(ast.NodeVisitor): 1abcdefghABCDijklmnopqrEFGHMNOPQRSTUVstuvwxyzIJKL

12 def __init__(self) -> None: 1abcdefghABCDijklmnopqrEFGHMNOPQRSTUVstuvwxyzIJKL

13 super().__init__() 1abcdefghABCDijklmnopqrEFGHstuvwxyzIJKL

14 

15 self.target: str | None = None 1abcdefghABCDijklmnopqrEFGHstuvwxyzIJKL

16 self.attrs: dict[str, str] = {} 1abcdefghABCDijklmnopqrEFGHstuvwxyzIJKL

17 self.previous_node_type: type[ast.AST] | None = None 1abcdefghABCDijklmnopqrEFGHstuvwxyzIJKL

18 

19 def visit(self, node: ast.AST) -> Any: 1abcdefghABCDijklmnopqrEFGHMNOPQRSTUVstuvwxyzIJKL

20 node_result = super().visit(node) 1abcdefghABCDijklmnopqrEFGHstuvwxyzIJKL

21 self.previous_node_type = type(node) 1abcdefghABCDijklmnopqrEFGHstuvwxyzIJKL

22 return node_result 1abcdefghABCDijklmnopqrEFGHstuvwxyzIJKL

23 

24 def visit_AnnAssign(self, node: ast.AnnAssign) -> Any: 1abcdefghABCDijklmnopqrEFGHMNOPQRSTUVstuvwxyzIJKL

25 if isinstance(node.target, ast.Name): 25 ↛ exitline 25 didn't return from function 'visit_AnnAssign' because the condition on line 25 was always true1abcdefghABCDijklmnopqrEFGHstuvwxyzIJKL

26 self.target = node.target.id 1abcdefghABCDijklmnopqrEFGHstuvwxyzIJKL

27 

28 def visit_Expr(self, node: ast.Expr) -> Any: 1abcdefghABCDijklmnopqrEFGHMNOPQRSTUVstuvwxyzIJKL

29 if ( 1abcdefghijklmnopqrstuvwxyz

30 isinstance(node.value, ast.Constant) 

31 and isinstance(node.value.value, str) 

32 and self.previous_node_type is ast.AnnAssign 

33 ): 

34 docstring = inspect.cleandoc(node.value.value) 1abcdefghABCDijklmnopqrEFGHstuvwxyzIJKL

35 if self.target: 35 ↛ 37line 35 didn't jump to line 37 because the condition on line 35 was always true1abcdefghABCDijklmnopqrEFGHstuvwxyzIJKL

36 self.attrs[self.target] = docstring 1abcdefghABCDijklmnopqrEFGHstuvwxyzIJKL

37 self.target = None 1abcdefghABCDijklmnopqrEFGHstuvwxyzIJKL

38 

39 

40def _dedent_source_lines(source: list[str]) -> str: 1abcdefghABCDijklmnopqrEFGHMNOPQRSTUVstuvwxyzIJKL

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

42 dedent_source = textwrap.dedent(''.join(source)) 1abcdefghABCDijklmnopqrEFGHstuvwxyzIJKL

43 if dedent_source.startswith((' ', '\t')): 1abcdefghABCDijklmnopqrEFGHstuvwxyzIJKL

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

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

46 # in a function as a workaround. 

47 dedent_source = f'def dedent_workaround():\n{dedent_source}' 1abcdefghABCDijklmnopqrEFGHstuvwxyzIJKL

48 return dedent_source 1abcdefghABCDijklmnopqrEFGHstuvwxyzIJKL

49 

50 

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

52 frame = inspect.currentframe() 1abcdefghABCDijklmnopqrEFGHstuvwxyzIJKL

53 

54 while frame: 1abcdefghABCDijklmnopqrEFGHstuvwxyzIJKL

55 if inspect.getmodule(frame) is inspect.getmodule(cls): 1abcdefghABCDijklmnopqrEFGHstuvwxyzIJKL

56 lnum = frame.f_lineno 1abcdefghABCDijklmnopqrEFGHstuvwxyzIJKL

57 try: 1abcdefghABCDijklmnopqrEFGHstuvwxyzIJKL

58 lines, _ = inspect.findsource(frame) 1abcdefghABCDijklmnopqrEFGHstuvwxyzIJKL

59 except OSError: 

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

61 # we don't want to error here. 

62 pass 

63 else: 

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

65 dedent_source = _dedent_source_lines(block_lines) 1abcdefghABCDijklmnopqrEFGHstuvwxyzIJKL

66 try: 1abcdefghABCDijklmnopqrEFGHstuvwxyzIJKL

67 block_tree = ast.parse(dedent_source) 1abcdefghABCDijklmnopqrEFGHstuvwxyzIJKL

68 except SyntaxError: 

69 pass 

70 else: 

71 stmt = block_tree.body[0] 1abcdefghABCDijklmnopqrEFGHstuvwxyzIJKL

72 if isinstance(stmt, ast.FunctionDef) and stmt.name == 'dedent_workaround': 1abcdefghABCDijklmnopqrEFGHstuvwxyzIJKL

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

74 stmt = stmt.body[0] 1abcdefghABCDijklmnopqrEFGHstuvwxyzIJKL

75 if isinstance(stmt, ast.ClassDef) and stmt.name == cls.__name__: 1abcdefghABCDijklmnopqrEFGHstuvwxyzIJKL

76 return block_lines 1abcdefghABCDijklmnopqrEFGHstuvwxyzIJKL

77 

78 frame = frame.f_back 1abcdefghABCDijklmnopqrEFGHstuvwxyzIJKL

79 

80 

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

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

83 

84 Args: 

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

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

87 the `inspect` module instead. 

88 

89 Returns: 

90 A mapping containing attribute names and their corresponding docstring. 

91 """ 

92 if use_inspect: 1abcdefghABCDijklmnopqrEFGHstuvwxyzIJKL

93 # Might not work as expected if two classes have the same name in the same source file. 

94 try: 1abcdefghABCDijklmnopqrEFGHstuvwxyzIJKL

95 source, _ = inspect.getsourcelines(cls) 1abcdefghABCDijklmnopqrEFGHstuvwxyzIJKL

96 except OSError: 

97 return {} 

98 else: 

99 source = _extract_source_from_frame(cls) 1abcdefghABCDijklmnopqrEFGHstuvwxyzIJKL

100 

101 if not source: 1abcdefghABCDijklmnopqrEFGHstuvwxyzIJKL

102 return {} 1abcdefghABCDijklmnopqrEFGHstuvwxyzIJKL

103 

104 dedent_source = _dedent_source_lines(source) 1abcdefghABCDijklmnopqrEFGHstuvwxyzIJKL

105 

106 visitor = DocstringVisitor() 1abcdefghABCDijklmnopqrEFGHstuvwxyzIJKL

107 visitor.visit(ast.parse(dedent_source)) 1abcdefghABCDijklmnopqrEFGHstuvwxyzIJKL

108 return visitor.attrs 1abcdefghABCDijklmnopqrEFGHstuvwxyzIJKL