diff --git a/doc_in_py/core.py b/doc_in_py/core.py index a7048ee..c0eaf77 100644 --- a/doc_in_py/core.py +++ b/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): diff --git a/doc_in_py/test/testing_module_multi/.tree.yml b/doc_in_py/test/testing_module_multi/.tree.yml index abc02bd..ff45c3a 100644 --- a/doc_in_py/test/testing_module_multi/.tree.yml +++ b/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 \ No newline at end of file