116 lines
4.1 KiB
Python
116 lines
4.1 KiB
Python
"""
|
|
* SPDX-FileCopyrightText: Copyright (c) 2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
|
|
* SPDX-License-Identifier: Apache-2.0
|
|
*
|
|
* Licensed under the Apache License, Version 2.0 (the "License");
|
|
* you may not use this file except in compliance with the License.
|
|
* You may obtain a copy of the License at
|
|
*
|
|
* https://www.apache.org/licenses/LICENSE-2.0
|
|
*
|
|
* Unless required by applicable law or agreed to in writing, software
|
|
* distributed under the License is distributed on an "AS IS" BASIS,
|
|
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
|
* See the License for the specific language governing permissions and
|
|
* limitations under the License.
|
|
"""
|
|
import os
|
|
import re
|
|
import subprocess
|
|
import sys
|
|
import textwrap
|
|
|
|
|
|
def update_comfyui_rtx_remix_readme(module, readme_path, section_header):
|
|
"""Generate docs for node types"""
|
|
section_header_style = section_header.split(' ')[0]
|
|
header_num = section_header_style.count('#')
|
|
|
|
output = []
|
|
nodes_by_module = {}
|
|
for name, node in module.NODE_CLASS_MAPPINGS.items():
|
|
nodes_by_module.setdefault(node.__module__, []).append((name, node))
|
|
|
|
for node_module, nodes in nodes_by_module.items():
|
|
output.append(f"{'#' * (header_num + 1)} {node_module.split('.')[-1].capitalize()}\n")
|
|
for node_info in nodes:
|
|
name, node_class = node_info
|
|
display_name = module.NODE_DISPLAY_NAME_MAPPINGS[name]
|
|
if node_class.__doc__:
|
|
# first line only
|
|
doc = node_class.__doc__.strip("\n")
|
|
line_break_index = doc.find("\n")
|
|
if line_break_index > 0:
|
|
doc = doc[:]
|
|
output.append(f"- **{display_name}**: {doc}\n")
|
|
else:
|
|
output.append(f"**{display_name}**\n")
|
|
|
|
output.append("\n")
|
|
|
|
replace_section(readme_path, section_header, output)
|
|
|
|
|
|
def replace_section(readme_path, section_header, contents):
|
|
"""Replace a section under a specific markdown header with new contents"""
|
|
section_header_style = section_header.split(' ')[0]
|
|
with open(readme_path, 'r') as f:
|
|
data = f.readlines()
|
|
|
|
new = []
|
|
start = -1
|
|
end = None
|
|
found = False
|
|
for i, line in enumerate(data):
|
|
# search for the first occurence of the special section header string
|
|
if not found:
|
|
if line.strip() == section_header:
|
|
found = True
|
|
start = i + 1
|
|
# search for the next header of the same magnitude
|
|
elif re.match(f"{section_header_style}[^#].*", line):
|
|
end = i
|
|
break
|
|
|
|
final = data[:start] + contents
|
|
if end:
|
|
final += data[end:]
|
|
with open(readme_path, 'w') as f:
|
|
f.writelines(final)
|
|
|
|
|
|
def update_readme(module, readme_path, section_header):
|
|
"""Generic `class_docs` entry point. Add repo-specific implementation here!"""
|
|
update_comfyui_rtx_remix_readme(module, readme_path, section_header)
|
|
|
|
|
|
def setup_repo_tool(parser, _):
|
|
parser.prog = "class_docs"
|
|
parser.description = "Generate summary documentation for a bunch of classes"
|
|
|
|
def run_repo_tool(options, config):
|
|
import subprocess
|
|
settings = config["repo_class_docs"]
|
|
|
|
# Some projects will require using their own python in order to be able to import and inspect
|
|
# classes so we use a subprocess
|
|
python_code = textwrap.dedent(f"""\
|
|
import sys
|
|
sys.path.append("{os.path.dirname(__file__)}")
|
|
import class_docs
|
|
import importlib.util
|
|
module_name = "{settings["module_name"]}"
|
|
spec = importlib.util.spec_from_file_location(module_name, "{settings["module"]}")
|
|
module = importlib.util.module_from_spec(spec)
|
|
sys.modules[module_name] = module
|
|
spec.loader.exec_module(module)
|
|
|
|
class_docs.update_readme(module, "{settings["file_path"]}", "{settings["section_header"]}")
|
|
""")
|
|
python_cmd = settings["python_install"]
|
|
|
|
if subprocess.call([python_cmd, "-c", python_code]) == 0:
|
|
print("Success!")
|
|
|
|
return run_repo_tool
|