feat: support .tree.yml dividing modules in a folder

doc_in_py/core.py

@@ -4,7 +4,10 @@
 import inspect
 import os
 import pkgutil
-from typing import List
+from collections import defaultdict
+from typing import List, Dict
+
+import yaml
 
 from doc_in_py import Struct
 from doc_in_py.comment_parser import prepare_raw_comment_struct, parse_raw_comments
@@ -13,7 +16,7 @@
 def get_module_members(module) -> Struct:
     root_path = module.__file__.split("__init__.py")[0]
     module_struct, sub_modules = extract_module_tree_without_comment(module, root_path)
-
+    sub_modules: Dict[str, Struct] | List[Struct]
     # Check whether the module has no source code.
     # It will raise OSError if the module has no source code
     try:
@@ -40,15 +43,23 @@ def get_module_members(module) -> Struct:
 def extract_module_tree_without_comment(module, root_path):
     module_struct: Struct = Struct("module", module, None, module.__name__)
 
-    # Get the submodules
+    # Get the sub_modules
+    sub_modules = get_sub_modules(module)
+
+    # Get sub_module tree from .tree.yml if it exists
     module_dir = os.path.dirname(inspect.getfile(module))
-    sub_modules = []
-    is_pkg = hasattr(module, "__path__")
-    if is_pkg:
-        for sub_module_info in pkgutil.iter_modules([module_dir]):
-            sub_module = importlib.import_module(
-                module.__name__ + "." + sub_module_info.name)
-            sub_modules.append(sub_module)
+    tree_yml_path = os.path.join(module_dir, ".tree.yml")
+    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"):
+        sections_dict = build_module_section_dict(tree_config_dict, sub_modules)
+        if sections_dict is not None:
+            sub_modules = sections_dict
 
     # Get the classes and functions
     true_members = []
@@ -78,7 +89,8 @@ def extract_module_tree_without_comment(module, root_path):
         # check decorators
         if hasattr(member, "__docinpy_todo"):
             member_type = "function" if inspect.isfunction(member) else "class"
-            todo_struct = Struct("todo", member.__docinpy_todo, pos, f"{member_type}: {member.__name__}")
+            todo_struct = Struct("todo", member.__docinpy_todo, pos,
+                                 f"{member_type}: {member.__name__}")
             n_todo += 1
             module_struct.children.append(todo_struct)
             parent_struct = todo_struct
@@ -92,6 +104,58 @@ def extract_module_tree_without_comment(module, root_path):
     return module_struct, sub_modules
 
 
+def build_module_section_dict(tree_config_dict, sub_modules):
+    if "sections" not in tree_config_dict:
+        return None
+    sections_dict = tree_config_dict["sections"]
+    if len(sections_dict) == 0:
+        return
+    sections_dict: Dict
+    name_to_module = {}
+    for sub_module in sub_modules:
+        name = sub_module.__name__.split(".")[-1]
+        name_to_module[name] = sub_module
+    map_module_name_to_module(sections_dict, name_to_module)
+    if len(name_to_module) > 0:
+        default_section_title = tree_config_dict.get("default section", "Others")
+        sections_dict[default_section_title] = list(name_to_module.values())
+    return sections_dict
+
+
+def get_sub_modules(module):
+    module_dir = os.path.dirname(inspect.getfile(module))
+    sub_modules = []
+    is_pkg = hasattr(module, "__path__")
+    if is_pkg:
+        for sub_module_info in pkgutil.iter_modules([module_dir]):
+            sub_module = importlib.import_module(
+                module.__name__ + "." + sub_module_info.name)
+            sub_modules.append(sub_module)
+    return sub_modules
+
+def map_module_name_to_module(sections_dict: Dict, name_to_module: Dict):
+    """
+    Recursively map the module name to module. Delete the module name from name_to_module in the process.
+    The module that is not map remains in name_to_module.
+    In the result, the leaf elements can be module, str, or None.
+    The str elements are the docstrings of the section
+    The None elements are the module not found
+    """
+    for title, content_list in sections_dict.items():
+        for i, content in enumerate(content_list):
+            if isinstance(content, str):
+                if content in name_to_module:
+                    content_list[i] = name_to_module[content]
+                    del name_to_module[content]
+                else:
+                    if i == 0:
+                        #content_list[i] = content_list[i]
+                        pass
+                    else:
+                        content_list[i] = None
+            elif isinstance(content, dict):
+                map_module_name_to_module(content, name_to_module)
+
 def add_function_class_to_struct(member, parent_struct, name, pos):
     if inspect.isclass(member):
         class_struct = Struct("class", member, pos, name)
@@ -168,9 +232,27 @@ def build_section_tree(root_struct: Struct):
 
 
 def process_sub_modules(sub_modules, root_struct: Struct):
-    for i, sub_module in enumerate(sub_modules):
-        member = sub_module
-        root_struct.children.append(get_module_members(member))
+    if isinstance(sub_modules, list):
+        for i, sub_module in enumerate(sub_modules):
+            if isinstance(sub_module, dict):
+                process_sub_modules(sub_module, root_struct)
+            elif isinstance(sub_module, str):
+                root_struct.children.append(Struct("comment", sub_module, None, None))
+            elif sub_module is not None:
+                root_struct.children.append(get_module_members(sub_module))
+    elif isinstance(sub_modules, dict):
+        for title, sub_module_list in sub_modules.items():
+            if root_struct.struct_type == "section":
+                new_level = root_struct.obj[1] + 1
+                new_section = Struct("section", (title, new_level), None, title)
+                process_sub_modules(sub_module_list, new_section)
+            else:
+                new_section = Struct("section", (title, 1), None, title)
+                process_sub_modules(sub_module_list, new_section)
+            root_struct.children.append(new_section)
+    else:
+        assert False
+
 
 
 def add_raw_comments_to_struct(cmt_structs: List[Struct], root_struct: Struct):

doc_in_py/test/testing_module_multi/.tree.yml

@@ -1,10 +1,12 @@
 sections:
   first section:
+    - This is some annotation to the first section
     - a
     - b
   second section:
     - c
     - another section:
+      - This is another section
       - d
       - e
 default section: other section