Coverage for /home/tbone/.local/share/hatch/env/virtual/importnb-aVRh-lqt/test.stdlib/lib/python3.9/site-packages/importnb/docstrings.py: 95%

37 statements  

« prev     ^ index     » next       coverage.py v6.5.0, created at 2022-10-03 09:31 -0700

1# coding: utf-8 

2"""# Special handling of markdown cells as docstrings. 

3 

4Modify the Python `ast` to assign docstrings to functions when they are preceded by a Markdown cell. 

5""" 

6 

7import ast 

8 

9"""# Modifying the `ast` 

10 

11 >>> assert isinstance(create_test, ast.Assign) 

12 >>> assert isinstance(test_update, ast.Attribute) 

13""" 

14 

15create_test = ast.parse( 

16 """__test__ = globals().get('__test__', {})""", mode="single" 

17).body[0] 

18test_update = ast.parse("""__test__.update""", mode="single").body[0].value 

19str_nodes = (ast.Str,) 

20 

21"""`TestStrings` is an `ast.NodeTransformer` that captures `str_nodes` in the `TestStrings.strings` object. 

22 

23```ipython 

24>>> assert isinstance(ast.parse(TestStrings().visit(ast.parse('"Test me"'))), ast.Module) 

25 

26``` 

27""" 

28 

29 

30class TestStrings(ast.NodeTransformer): 

31 

32 strings = None 

33 

34 def visit_Module(self, module): 

35 """`TestStrings.visit_Module` initializes the capture. After all the nodes are visit we append `create_test and test_update` 

36 to populate the `"__test__"` attribute. 

37 """ 

38 self.strings = [] 

39 module = self.visit_body(module) 

40 module.body += ( 

41 [create_test] 

42 + [ 

43 ast.copy_location( 

44 ast.Expr( 

45 ast.Call( 

46 func=test_update, 

47 args=[ 

48 ast.Dict( 

49 keys=[ast.Str("string-{}".format(node.lineno))], 

50 values=[node], 

51 ) 

52 ], 

53 keywords=[], 

54 ) 

55 ), 

56 node, 

57 ) 

58 for node in self.strings 

59 ] 

60 if self.strings 

61 else [] 

62 ) 

63 return module 

64 

65 def visit_body(self, node): 

66 """`TestStrings.visit_body` visits nodes with a `"body"` attibute and extracts potential string tests.""" 

67 

68 body = [] 

69 if ( 

70 node.body 

71 and isinstance(node.body[0], ast.Expr) 

72 and isinstance(node.body[0].value, str_nodes) 

73 ): 

74 body.append(node.body.pop(0)) 

75 node.body = body + [ 

76 (self.visit_body if hasattr(object, "body") else self.visit)(object) 

77 for object in node.body 

78 ] 

79 return node 

80 

81 def visit_Expr(self, node): 

82 """`TestStrings.visit_Expr` append the `str_nodes` to `TestStrings.strings` to append to the `ast.Module`.""" 

83 

84 if isinstance(node.value, str_nodes): 

85 self.strings.append( 

86 ast.copy_location(ast.Str(node.value.s.replace("\n```", "\n")), node) 

87 ) 

88 return node 

89 

90 

91def update_docstring(module): 

92 from functools import reduce 

93 

94 module.body = reduce(markdown_docstring, module.body, []) 

95 return TestStrings().visit(module) 

96 

97 

98docstring_ast_types = ast.ClassDef, ast.FunctionDef 

99try: 

100 docstring_ast_types += (ast.AsyncFunctionDef,) 

101except: 

102 ... 

103 

104 

105def markdown_docstring(nodes, node): 

106 if ( 

107 len(nodes) > 1 

108 and str_expr(nodes[-1]) 

109 and isinstance(node, docstring_ast_types) 

110 and not str_expr(node.body[0]) 

111 ): 

112 node.body.insert(0, nodes.pop()) 

113 return nodes.append(node) or nodes 

114 

115 

116def str_expr(node): 

117 return isinstance(node, ast.Expr) and isinstance(node.value, ast.Str)