test: add test case for .tree.yml

EvoEvolver · Oct 9, 2023 · 68160cb · 68160cb
1 parent 3aa39cd
commit 68160cb
Show file tree

Hide file tree

Showing 2 changed files with 27 additions and 1 deletion.
diff --git a/doc_in_py/core.py b/doc_in_py/core.py
@@ -49,14 +49,15 @@ def extract_module_tree_without_comment(module, root_path):
     # Get sub_module tree from .tree.yml if it exists
     module_dir = os.path.dirname(inspect.getfile(module))
     tree_yml_path = os.path.join(module_dir, ".tree.yml")
+    tree_config_dict = {}
     if os.path.exists(tree_yml_path):
         try:
             with open(tree_yml_path, "r") as f:
                 tree_config_dict = yaml.load(f, Loader=yaml.FullLoader)
         except:
             pass
 
-    if module.__file__.endswith("__init__.py"):
+    if module.__file__.endswith("__init__.py") and len(tree_config_dict) > 0:
         sections_dict = build_module_section_dict(tree_config_dict, sub_modules)
         if sections_dict is not None:
             sub_modules = sections_dict

diff --git a/doc_in_py/test/test_module_multi.py b/doc_in_py/test/test_module_multi.py
@@ -0,0 +1,25 @@
+from doc_in_py.core import get_module_members
+from utils import to_dict
+
+
+def test_section():
+    import testing_module_multi
+    struct = get_module_members(testing_module_multi)
+    actual = to_dict(struct)
+    expected = {'struct_type': 'module', 'name': 'testing_module_multi', 'children': [
+        {'struct_type': 'section', 'name': 'first section', 'children': [
+            {'struct_type': 'comment',
+             'obj': 'This is some annotation to the first section'},
+            {'struct_type': 'module', 'name': 'testing_module_multi.a',
+             'children': [{'struct_type': 'class', 'name': 'A'}]},
+            {'struct_type': 'module', 'name': 'testing_module_multi.b'}]},
+        {'struct_type': 'section', 'name': 'second section',
+         'children': [{'struct_type': 'module', 'name': 'testing_module_multi.c'},
+                      {'struct_type': 'section', 'name': 'another section', 'children': [
+                          {'struct_type': 'comment', 'obj': 'This is another section'},
+                          {'struct_type': 'module', 'name': 'testing_module_multi.d'},
+                          {'struct_type': 'module', 'name': 'testing_module_multi.e'}]}]},
+        {'struct_type': 'section', 'name': 'other section', 'children': [
+            {'struct_type': 'module', 'name': 'testing_module_multi.another'}]}]}
+
+    assert actual == expected