Coverage for /home/tbone/.local/share/hatch/env/virtual/importnb-aVRh-lqt/test.interactive/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
« 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.
4Modify the Python `ast` to assign docstrings to functions when they are preceded by a Markdown cell.
5"""
7import ast
9"""# Modifying the `ast`
11 >>> assert isinstance(create_test, ast.Assign)
12 >>> assert isinstance(test_update, ast.Attribute)
13"""
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,)
21"""`TestStrings` is an `ast.NodeTransformer` that captures `str_nodes` in the `TestStrings.strings` object.
23```ipython
24>>> assert isinstance(ast.parse(TestStrings().visit(ast.parse('"Test me"'))), ast.Module)
26```
27"""
30class TestStrings(ast.NodeTransformer):
32 strings = None
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
65 def visit_body(self, node):
66 """`TestStrings.visit_body` visits nodes with a `"body"` attibute and extracts potential string tests."""
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
81 def visit_Expr(self, node):
82 """`TestStrings.visit_Expr` append the `str_nodes` to `TestStrings.strings` to append to the `ast.Module`."""
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
91def update_docstring(module):
92 from functools import reduce
94 module.body = reduce(markdown_docstring, module.body, [])
95 return TestStrings().visit(module)
98docstring_ast_types = ast.ClassDef, ast.FunctionDef
99try:
100 docstring_ast_types += (ast.AsyncFunctionDef,)
101except:
102 ...
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
116def str_expr(node):
117 return isinstance(node, ast.Expr) and isinstance(node.value, ast.Str)