122 lines
6.0 KiB
Python
122 lines
6.0 KiB
Python
import os
|
||
from textwrap import indent
|
||
|
||
class FileNode:
|
||
def __init__(self, name):
|
||
self.name = name
|
||
self.children = []
|
||
self.is_leaf = False
|
||
self.level = 0
|
||
self.parenting_ship = []
|
||
self.comment = ""
|
||
self.comment_maxlen_show = 50
|
||
|
||
@staticmethod
|
||
def add_linebreaks_at_spaces(string, interval=10):
|
||
return '\n'.join(string[i:i+interval] for i in range(0, len(string), interval))
|
||
|
||
def sanitize_comment(self, comment):
|
||
if len(comment) > self.comment_maxlen_show: suf = '...'
|
||
else: suf = ''
|
||
comment = comment[:self.comment_maxlen_show]
|
||
comment = comment.replace('\"', '').replace('`', '').replace('\n', '').replace('`', '').replace('$', '')
|
||
comment = self.add_linebreaks_at_spaces(comment, 10)
|
||
return '`' + comment + suf + '`'
|
||
|
||
def add_file(self, file_path, file_comment):
|
||
directory_names, file_name = os.path.split(file_path)
|
||
current_node = self
|
||
level = 1
|
||
if directory_names == "":
|
||
new_node = FileNode(file_name)
|
||
current_node.children.append(new_node)
|
||
new_node.is_leaf = True
|
||
new_node.comment = self.sanitize_comment(file_comment)
|
||
new_node.level = level
|
||
current_node = new_node
|
||
else:
|
||
dnamesplit = directory_names.split(os.sep)
|
||
for i, directory_name in enumerate(dnamesplit):
|
||
found_child = False
|
||
level += 1
|
||
for child in current_node.children:
|
||
if child.name == directory_name:
|
||
current_node = child
|
||
found_child = True
|
||
break
|
||
if not found_child:
|
||
new_node = FileNode(directory_name)
|
||
current_node.children.append(new_node)
|
||
new_node.level = level - 1
|
||
current_node = new_node
|
||
term = FileNode(file_name)
|
||
term.level = level
|
||
term.comment = self.sanitize_comment(file_comment)
|
||
term.is_leaf = True
|
||
current_node.children.append(term)
|
||
|
||
def print_files_recursively(self, level=0, code="R0"):
|
||
print(' '*level + self.name + ' ' + str(self.is_leaf) + ' ' + str(self.level))
|
||
for j, child in enumerate(self.children):
|
||
child.print_files_recursively(level=level+1, code=code+str(j))
|
||
self.parenting_ship.extend(child.parenting_ship)
|
||
p1 = f"""{code}[\"🗎{self.name}\"]""" if self.is_leaf else f"""{code}[[\"📁{self.name}\"]]"""
|
||
p2 = """ --> """
|
||
p3 = f"""{code+str(j)}[\"🗎{child.name}\"]""" if child.is_leaf else f"""{code+str(j)}[[\"📁{child.name}\"]]"""
|
||
edge_code = p1 + p2 + p3
|
||
if edge_code in self.parenting_ship:
|
||
continue
|
||
self.parenting_ship.append(edge_code)
|
||
if self.comment != "":
|
||
pc1 = f"""{code}[\"🗎{self.name}\"]""" if self.is_leaf else f"""{code}[[\"📁{self.name}\"]]"""
|
||
pc2 = f""" -.-x """
|
||
pc3 = f"""C{code}[\"{self.comment}\"]:::Comment"""
|
||
edge_code = pc1 + pc2 + pc3
|
||
self.parenting_ship.append(edge_code)
|
||
|
||
|
||
MERMAID_TEMPLATE = r"""
|
||
```mermaid
|
||
flowchart LR
|
||
%% <gpt_academic_hide_mermaid_code> 一个特殊标记，用于在生成mermaid图表时隐藏代码块
|
||
classDef Comment stroke-dasharray: 5 5
|
||
subgraph {graph_name}
|
||
{relationship}
|
||
end
|
||
```
|
||
"""
|
||
|
||
def build_file_tree_mermaid_diagram(file_manifest, file_comments, graph_name):
|
||
# Create the root node
|
||
file_tree_struct = FileNode("root")
|
||
# Build the tree structure
|
||
for file_path, file_comment in zip(file_manifest, file_comments):
|
||
file_tree_struct.add_file(file_path, file_comment)
|
||
file_tree_struct.print_files_recursively()
|
||
cc = "\n".join(file_tree_struct.parenting_ship)
|
||
ccc = indent(cc, prefix=" "*8)
|
||
return MERMAID_TEMPLATE.format(graph_name=graph_name, relationship=ccc)
|
||
|
||
if __name__ == "__main__":
|
||
# File manifest
|
||
file_manifest = [
|
||
"cradle_void_terminal.ipynb",
|
||
"tests/test_utils.py",
|
||
"tests/test_plugins.py",
|
||
"tests/test_llms.py",
|
||
"config.py",
|
||
"build/ChatGLM-6b-onnx-u8s8/chatglm-6b-int8-onnx-merged/model_weights_0.bin",
|
||
"crazy_functions/latex_fns/latex_actions.py",
|
||
"crazy_functions/latex_fns/latex_toolbox.py"
|
||
]
|
||
file_comments = [
|
||
"根据位置和名称，可能是一个模块的初始化文件根据位置和名称，可能是一个模块的初始化文件根据位置和名称，可能是一个模块的初始化文件",
|
||
"包含一些用于文本处理和模型微调的函数和装饰器包含一些用于文本处理和模型微调的函数和装饰器包含一些用于文本处理和模型微调的函数和装饰器",
|
||
"用于构建HTML报告的类和方法用于构建HTML报告的类和方法用于构建HTML报告的类和方法",
|
||
"包含了用于文本切分的函数，以及处理PDF文件的示例代码包含了用于文本切分的函数，以及处理PDF文件的示例代码包含了用于文本切分的函数，以及处理PDF文件的示例代码",
|
||
"用于解析和翻译PDF文件的功能和相关辅助函数用于解析和翻译PDF文件的功能和相关辅助函数用于解析和翻译PDF文件的功能和相关辅助函数",
|
||
"是一个包的初始化文件，用于初始化包的属性和导入模块是一个包的初始化文件，用于初始化包的属性和导入模块是一个包的初始化文件，用于初始化包的属性和导入模块",
|
||
"用于加载和分割文件中的文本的通用文件加载器用于加载和分割文件中的文本的通用文件加载器用于加载和分割文件中的文本的通用文件加载器",
|
||
"包含了用于构建和管理向量数据库的函数和类包含了用于构建和管理向量数据库的函数和类包含了用于构建和管理向量数据库的函数和类",
|
||
]
|
||
print(build_file_tree_mermaid_diagram(file_manifest, file_comments, "项目文件树")) |