Skip to content

GitLab

Switch branch/tag
Find file Normal viewHistoryPermalink
auto_generate_api.py 3.32 KB
Newer Older
mashun1's avatar
controlnet  
mashun1 committed
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97 
import os
from os import path as osp


def scandir(dir_path, suffix=None, recursive=False, full_path=False):
    """Scan a directory to find the interested files.

    Args:
        dir_path (str): Path of the directory.
        suffix (str | tuple(str), optional): File suffix that we are
            interested in. Default: None.
        recursive (bool, optional): If set to True, recursively scan the
            directory. Default: False.
        full_path (bool, optional): If set to True, include the dir_path.
            Default: False.

    Returns:
        A generator for all the interested files with relative paths.
    """

    if (suffix is not None) and not isinstance(suffix, (str, tuple)):
        raise TypeError('"suffix" must be a string or tuple of strings')

    root = dir_path

    def _scandir(dir_path, suffix, recursive):
        for entry in os.scandir(dir_path):
            if not entry.name.startswith('.') and entry.is_file():
                if full_path:
                    return_path = entry.path
                else:
                    return_path = osp.relpath(entry.path, root)

                if suffix is None:
                    yield return_path
                elif return_path.endswith(suffix):
                    yield return_path
            else:
                if recursive:
                    yield from _scandir(entry.path, suffix=suffix, recursive=recursive)
                else:
                    continue

    return _scandir(dir_path, suffix=suffix, recursive=recursive)


# specifically generate a fake __init__.py for scripts folder to generate docs
with open('../scripts/__init__.py', 'w') as f:
    pass

module_name_list = ['basicsr', 'scripts']
for module_name in module_name_list:
    cur_dir = osp.abspath(osp.dirname(__file__))
    output_dir = osp.join(cur_dir, 'api')
    module_dir = osp.join(osp.dirname(cur_dir), module_name)
    os.makedirs(output_dir, exist_ok=True)

    api_content = f'{module_name} API\n=========================\n'
    submodule_name_list = []

    for path in sorted(scandir(module_dir, suffix='.py', recursive=True)):
        if path in ['__init__.py', 'version.py']:
            continue

        path = f'{module_name}.' + path.replace('\\', '/').replace('/', '.').replace('.py', '.rst')

        # create .rst file
        output_rst = osp.join(output_dir, path)
        with open(output_rst, 'w') as f:
            # write contents
            title = path.replace('.rst', '').replace('_', '\\_')
            content = f'{title}\n===================================================================================\n'
            automodule = path.replace('.rst', '')
            content += f'\n.. automodule:: {automodule}'
            content += r'''
    :members:
    :undoc-members:
    :show-inheritance:
    '''
            f.write(content)

        # add to api.rst
        submodule_name = path.split('.')[1]
        if submodule_name not in submodule_name_list:
            submodule_name_list.append(submodule_name)
            api_content += f'\n\n{module_name}.{submodule_name}\n-----------------------------------------------------'
            api_content += r'''
.. toctree::
    :maxdepth: 4

'''

        api_content += f'    {automodule}\n'

    # write to api.rst
    with open(os.path.join(output_dir, f'api_{module_name}.rst'), 'w') as f:
        f.write(api_content)