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

1# coding: utf-8

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

3

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

5"""

6

7 import ast

8

9"""# Modifying the `ast`

10

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

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

13"""

14

15 create_test = ast.parse(

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

17).body[0]

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

19 str_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

30 class 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

91 def update_docstring(module):

92 from functools import reduce

93

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

95 return TestStrings().visit(module)

96

97

98 docstring_ast_types = ast.ClassDef, ast.FunctionDef

99 try:

100 docstring_ast_types += (ast.AsyncFunctionDef,)

101 except:

102 ...

103

104

105 def 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

116 def str_expr(node):

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