129 lines
5.0 KiB
Python
129 lines
5.0 KiB
Python
# Ultralytics YOLO 🚀, AGPL-3.0 license
|
|
"""
|
|
Helper file to build Ultralytics Docs reference section. Recursively walks through ultralytics dir and builds an MkDocs
|
|
reference section of *.md files composed of classes and functions, and also creates a nav menu for use in mkdocs.yaml.
|
|
|
|
Note: Must be run from repository root directory. Do not run from docs directory.
|
|
"""
|
|
|
|
import re
|
|
from collections import defaultdict
|
|
from pathlib import Path
|
|
|
|
from ultralytics.utils import ROOT
|
|
|
|
NEW_YAML_DIR = ROOT.parent
|
|
CODE_DIR = ROOT
|
|
REFERENCE_DIR = ROOT.parent / 'docs/en/reference'
|
|
|
|
|
|
def extract_classes_and_functions(filepath: Path) -> tuple:
|
|
"""Extracts class and function names from a given Python file."""
|
|
content = filepath.read_text()
|
|
class_pattern = r'(?:^|\n)class\s(\w+)(?:\(|:)'
|
|
func_pattern = r'(?:^|\n)def\s(\w+)\('
|
|
|
|
classes = re.findall(class_pattern, content)
|
|
functions = re.findall(func_pattern, content)
|
|
|
|
return classes, functions
|
|
|
|
|
|
def create_markdown(py_filepath: Path, module_path: str, classes: list, functions: list):
|
|
"""Creates a Markdown file containing the API reference for the given Python module."""
|
|
md_filepath = py_filepath.with_suffix('.md')
|
|
|
|
# Read existing content and keep header content between first two ---
|
|
header_content = ''
|
|
if md_filepath.exists():
|
|
existing_content = md_filepath.read_text()
|
|
header_parts = existing_content.split('---')
|
|
for part in header_parts:
|
|
if 'description:' in part or 'comments:' in part:
|
|
header_content += f'---{part}---\n\n'
|
|
|
|
module_name = module_path.replace('.__init__', '')
|
|
module_path = module_path.replace('.', '/')
|
|
url = f'https://github.com/ultralytics/ultralytics/blob/main/{module_path}.py'
|
|
edit = f'https://github.com/ultralytics/ultralytics/edit/main/{module_path}.py'
|
|
title_content = (
|
|
f'# Reference for `{module_path}.py`\n\n'
|
|
f'!!! Note\n\n'
|
|
f' This file is available at [{url}]({url}). If you spot a problem please help fix it by [contributing](https://docs.ultralytics.com/help/contributing/) a [Pull Request]({edit}) 🛠️. Thank you 🙏!\n\n'
|
|
)
|
|
md_content = ['<br><br>\n'] + [f'## ::: {module_name}.{class_name}\n\n<br><br>\n' for class_name in classes]
|
|
md_content.extend(f'## ::: {module_name}.{func_name}\n\n<br><br>\n' for func_name in functions)
|
|
md_content = header_content + title_content + '\n'.join(md_content)
|
|
if not md_content.endswith('\n'):
|
|
md_content += '\n'
|
|
|
|
md_filepath.parent.mkdir(parents=True, exist_ok=True)
|
|
md_filepath.write_text(md_content)
|
|
|
|
return md_filepath.relative_to(NEW_YAML_DIR)
|
|
|
|
|
|
def nested_dict() -> defaultdict:
|
|
"""Creates and returns a nested defaultdict."""
|
|
return defaultdict(nested_dict)
|
|
|
|
|
|
def sort_nested_dict(d: dict) -> dict:
|
|
"""Sorts a nested dictionary recursively."""
|
|
return {key: sort_nested_dict(value) if isinstance(value, dict) else value for key, value in sorted(d.items())}
|
|
|
|
|
|
def create_nav_menu_yaml(nav_items: list):
|
|
"""Creates a YAML file for the navigation menu based on the provided list of items."""
|
|
nav_tree = nested_dict()
|
|
|
|
for item_str in nav_items:
|
|
item = Path(item_str)
|
|
parts = item.parts
|
|
current_level = nav_tree['reference']
|
|
for part in parts[2:-1]: # skip the first two parts (docs and reference) and the last part (filename)
|
|
current_level = current_level[part]
|
|
|
|
md_file_name = parts[-1].replace('.md', '')
|
|
current_level[md_file_name] = item
|
|
|
|
nav_tree_sorted = sort_nested_dict(nav_tree)
|
|
|
|
def _dict_to_yaml(d, level=0):
|
|
"""Converts a nested dictionary to a YAML-formatted string with indentation."""
|
|
yaml_str = ''
|
|
indent = ' ' * level
|
|
for k, v in d.items():
|
|
if isinstance(v, dict):
|
|
yaml_str += f'{indent}- {k}:\n{_dict_to_yaml(v, level + 1)}'
|
|
else:
|
|
yaml_str += f"{indent}- {k}: {str(v).replace('docs/en/', '')}\n"
|
|
return yaml_str
|
|
|
|
# Print updated YAML reference section
|
|
print('Scan complete, new mkdocs.yaml reference section is:\n\n', _dict_to_yaml(nav_tree_sorted))
|
|
|
|
# Save new YAML reference section
|
|
# (NEW_YAML_DIR / 'nav_menu_updated.yml').write_text(_dict_to_yaml(nav_tree_sorted))
|
|
|
|
|
|
def main():
|
|
"""Main function to extract class and function names, create Markdown files, and generate a YAML navigation menu."""
|
|
nav_items = []
|
|
|
|
for py_filepath in CODE_DIR.rglob('*.py'):
|
|
classes, functions = extract_classes_and_functions(py_filepath)
|
|
|
|
if classes or functions:
|
|
py_filepath_rel = py_filepath.relative_to(CODE_DIR)
|
|
md_filepath = REFERENCE_DIR / py_filepath_rel
|
|
module_path = f"ultralytics.{py_filepath_rel.with_suffix('').as_posix().replace('/', '.')}"
|
|
md_rel_filepath = create_markdown(md_filepath, module_path, classes, functions)
|
|
nav_items.append(str(md_rel_filepath))
|
|
|
|
create_nav_menu_yaml(nav_items)
|
|
|
|
|
|
if __name__ == '__main__':
|
|
main()
|