Coverage for src/pytkdocs/parsers/attributes.py: 95.98%

1"""Module containing functions to parse attributes in the source code.""" 

2 

3import ast 

4import inspect 

5from functools import lru_cache 

6from textwrap import dedent 

7from typing import get_type_hints 

8 

9try: 

10 from ast import unparse # type: ignore 

11except ImportError: 

12 from astunparse import unparse # type: ignore 

13 

14RECURSIVE_NODES = (ast.If, ast.IfExp, ast.Try, ast.With) 

15 

16 

17def get_nodes(obj): 

18 try: 

19 source = inspect.getsource(obj) 

20 except (OSError, TypeError): 

21 source = "" 

22 return ast.parse(dedent(source)).body 

23 

24 

25def recurse_on_node(node): 

26 if isinstance(node, ast.Try): 

27 yield from get_pairs(node.body) 

28 for handler in node.handlers: 

29 yield from get_pairs(handler.body) 

30 yield from get_pairs(node.orelse) 

31 yield from get_pairs(node.finalbody) 

32 elif isinstance(node, ast.If): 32 ↛ 36line 32 didn't jump to line 36, because the condition on line 32 was never false

33 yield from get_pairs(node.body) 

34 yield from get_pairs(node.orelse) 

35 else: 

36 yield from get_pairs(node.body) 

37 

38 

39def get_pairs(nodes): 

40 if len(nodes) < 2: 

41 return 

42 

43 index = 0 

44 while index < len(nodes): 

45 node1 = nodes[index] 

46 if index < len(nodes) - 1: 

47 node2 = nodes[index + 1] 

48 else: 

49 node2 = None 

50 if isinstance(node1, (ast.Assign, ast.AnnAssign)): 

51 if isinstance(node2, ast.Expr) and isinstance(node2.value, ast.Str): 

52 yield node1, node2.value 

53 index += 2 

54 else: 

55 yield node1, None 

56 index += 1 

57 else: 

58 index += 1 

59 if isinstance(node1, RECURSIVE_NODES): 

60 yield from recurse_on_node(node1) 

61 if isinstance(node2, RECURSIVE_NODES): 

62 yield from recurse_on_node(node2) 

63 index += 1 

64 elif not isinstance(node2, (ast.Assign, ast.AnnAssign)): 

65 index += 1 

66 

67 

68def get_module_or_class_attributes(nodes): 

69 result = {} 

70 for assignment, string_node in get_pairs(nodes): 

71 string = inspect.cleandoc(string_node.s) if string_node else None 

72 if isinstance(assignment, ast.Assign): 

73 names = [] 

74 for target in assignment.targets: 

75 if isinstance(target, ast.Name): 

76 names.append(target.id) 

77 elif isinstance(target, ast.Tuple): 

78 names.extend([name.id for name in target.elts]) 

79 else: 

80 names = [assignment.target.id] 

81 for name in names: 

82 result[name] = string 

83 return result 

84 

85 

86def combine(docstrings, type_hints): 

87 return { 

88 name: {"annotation": type_hints.get(name, inspect.Signature.empty), "docstring": docstrings.get(name)} 

89 for name in set(docstrings.keys()) | set(type_hints.keys()) 

90 } 

91 

92 

93def merge(base, extra): 

94 for attr_name, data in extra.items(): 

95 if attr_name in base: 

96 if data["annotation"] is not inspect.Signature.empty: 96 ↛ 98line 96 didn't jump to line 98, because the condition on line 96 was never false

97 base[attr_name]["annotation"] = data["annotation"] 

98 if data["docstring"] is not None: 98 ↛ 99line 98 didn't jump to line 99, because the condition on line 98 was never true

99 base[attr_name]["docstring"] = data["docstring"] 

100 else: 

101 base[attr_name] = data 

102 

103 

104@lru_cache() 

105def get_module_attributes(module): 

106 return combine(get_module_or_class_attributes(get_nodes(module)), get_type_hints(module)) 

107 

108 

109@lru_cache() 

110def get_class_attributes(cls): 

111 nodes = get_nodes(cls) 

112 if not nodes: 

113 return {} 

114 try: 

115 type_hints = get_type_hints(cls) 

116 except NameError: 

117 # The __config__ attribute (a class) of Pydantic models trigger this error: 

118 # NameError: name 'SchemaExtraCallable' is not defined 

119 type_hints = {} 

120 return combine(get_module_or_class_attributes(nodes[0].body), type_hints) 

121 

122 

123def pick_target(target): 

124 return isinstance(target, ast.Attribute) and isinstance(target.value, ast.Name) and target.value.id == "self" 

125 

126 

127def unparse_annotation(node): 

128 code = unparse(node).rstrip("\n") 

129 return code.replace("(", "").replace(")", "") 

130 

131 

132@lru_cache() 

133def get_instance_attributes(func): 

134 nodes = get_nodes(func) 

135 if not nodes: 

136 return {} 

137 

138 result = {} 

139 

140 for assignment, string in get_pairs(nodes[0].body): 

141 annotation = names = None 

142 if isinstance(assignment, ast.AnnAssign): 

143 if pick_target(assignment.target): 

144 names = [assignment.target.attr] 

145 annotation = unparse_annotation(assignment.annotation) 

146 else: 

147 names = [target.attr for target in assignment.targets if pick_target(target)] 

148 

149 if not names or (string is None and annotation is None): 

150 continue 

151 

152 docstring = inspect.cleandoc(string.s) if string else None 

153 for name in names: 

154 result[name] = {"annotation": annotation, "docstring": docstring} 

155 

156 return result