此方法适用于模块方法生成
import ast
import os
import platform
import matplotlib.pyplot as plt
import matplotlib.patches as patches
from matplotlib.patches import FancyBboxPatch
import textwrap
def get_system_fonts():
"""根据操作系统自动选择适配的中文字体（增加更多备选）"""
system = platform.system()
if system == "Darwin": # macOS系统，增加更多常见中文字体备选
return [
"PingFang SC", # 苹方（优先）
"Heiti TC", # 黑体
"Microsoft YaHei", # 微软雅黑（兼容）
"SimHei", # 黑体（兼容）
"Arial Unicode MS", # 系统默认Unicode字体
"Helvetica",
"Arial"
]
elif system == "Windows": # Windows系统
return [
"Microsoft YaHei", # 微软雅黑（优先）
"SimHei", # 黑体
"Arial"
]
else: # Linux或其他系统
return [
"WenQuanYi Micro Hei", # 文泉驿微米黑
"Heiti TC", # 黑体
"Arial"
]
def extract_classes_and_functions_from_file(filepath):
"""解析Python文件，提取类、类内方法和全局函数"""
with open(filepath, "r", encoding="utf-8") as file:
tree = ast.parse(file.read(), filename=filepath)
# 为所有节点添加parent属性，用于区分全局函数和类内方法
for node in ast.walk(tree):
for child in ast.iter_child_nodes(node):
child.parent = node
class_info = {} # 存储类名: [方法列表]
global_functions = [] # 存储全局函数
for node in ast.walk(tree):
# 提取类和类内方法
if isinstance(node, ast.ClassDef):
methods = [f"{subnode.name}()" for subnode in node.body if isinstance(subnode, ast.FunctionDef)]
class_info[node.name] = methods
# 提取全局函数（排除类内方法）
if isinstance(node, ast.FunctionDef) and not (hasattr(node, 'parent') and isinstance(node.parent, ast.ClassDef)):
global_functions.append(f"{node.name}()")
return class_info, global_functions
def plot_class_diagram(class_info, global_functions, output_path, module_name):
"""绘制IDEA风格类图（增强字体兼容性）"""
# 自动配置系统兼容字体
plt.rcParams["font.family"] = get_system_fonts()
plt.rcParams["axes.unicode_minus"] = False # 解决负号显示异常
# 动态计算画布高度，避免内容重叠
block_height = 0.35 # 每个类/函数块的高度
spacing = 0.15 # 块之间的间距
total_height = 0.6 # 顶部留白（预留模块名位置）
# 累加类所需高度（类名块 + 对应方法块）
for methods in class_info.values():
total_height += block_height + (block_height * len(methods)) + spacing
# 累加全局函数所需高度（标题块 + 对应函数块）
if global_functions:
total_height += block_height + (block_height * len(global_functions)) + spacing
total_height += 0.4 # 底部留白
# 创建画布（宽度固定10，高度根据内容动态调整）
fig, ax = plt.subplots(figsize=(10, max(6, total_height)))
ax.set_xlim(0, 1)
ax.set_ylim(0, total_height)
ax.set_facecolor('#ffffff') # 白色背景，贴近IDEA界面风格
ax.set_axis_off() # 隐藏默认坐标轴
current_y = total_height - 0.5 # 初始绘制位置（从顶部往下）
# 1. 绘制模块名标题（居中加粗，突出模块标识）
plt.text(0.5, current_y + 0.15, f"模块: {module_name}",
ha="center", va="center", fontsize=14, fontweight='bold', color='#333333')
current_y -= 0.3 # 下移，为类块预留位置
# 2. 绘制类（蓝色系，区分类名和方法）
for class_name, methods in class_info.items():
# 绘制类名块（深色边框+浅蓝色背景，IDEA类图风格）
class_rect = FancyBboxPatch(
(0.1, current_y - block_height), 0.8, block_height,
boxstyle="square,pad=0",
edgecolor='#2a76be', facecolor='#e6f2ff', lw=1.5 # 蓝色系配色
)
ax.add_patch(class_rect)
# 类名文本（左对齐，加粗）
plt.text(0.15, current_y - block_height/2, class_name,
ha="left", va="center", fontsize=12, fontweight='bold', color='#2a76be')
current_y -= block_height # 下移，为方法块预留位置
# 绘制当前类的方法块（浅色边框+更浅背景，与类名块区分）
for method in methods:
method_rect = FancyBboxPatch(
(0.1, current_y - block_height), 0.8, block_height,
boxstyle="square,pad=0",
edgecolor='#8ec0e4', facecolor='#f0f7ff', lw=1 # 浅蓝配色
)
ax.add_patch(method_rect)
# 方法名文本（左对齐，自动换行避免超长）
wrapped_method = textwrap.fill(method, width=45) # 控制每行长度
plt.text(0.15, current_y - block_height/2, wrapped_method,
ha="left", va="center", fontsize=10, color='#333333')
current_y -= block_height # 下移，为下一个方法/类预留位置
current_y -= spacing # 类之间的间距
# 3. 绘制全局函数（橙色系，与类区分）
if global_functions:
# 绘制全局函数标题块（深色边框+浅橙色背景）
func_title_rect = FancyBboxPatch(
(0.1, current_y - block_height), 0.8, block_height,
boxstyle="square,pad=0",
edgecolor='#e67e22', facecolor='#fff2e6', lw=1.5 # 橙色系配色
)
ax.add_patch(func_title_rect)
# 标题文本（左对齐，加粗）
plt.text(0.15, current_y - block_height/2, "全局函数",
ha="left", va="center", fontsize=12, fontweight='bold', color='#e67e22')
current_y -= block_height # 下移，为函数块预留位置
# 绘制全局函数块（浅色边框+更浅背景）
for func in global_functions:
func_rect = FancyBboxPatch(
(0.1, current_y - block_height), 0.8, block_height,
boxstyle="square,pad=0",
edgecolor='#f3b664', facecolor='#fff7f0', lw=1 # 浅橙配色
)
ax.add_patch(func_rect)
# 函数名文本（左对齐，自动换行）
wrapped_func = textwrap.fill(func, width=45)
plt.text(0.15, current_y - block_height/2, wrapped_func,
ha="left", va="center", fontsize=10, color='#333333')
current_y -= block_height # 下移，为下一个函数预留位置
current_y -= spacing # 与其他内容区隔
# 保存类图（高分辨率300dpi，避免内容截断）
plt.tight_layout()
plt.savefig(output_path, format="png", dpi=300, bbox_inches='tight', facecolor='#ffffff')
plt.close()
print(f"类图已保存：{output_path}")
def generate_class_diagram(input_path):
"""批量生成指定文件夹下所有Python文件的类图"""
# 校验输入路径有效性
if not os.path.isdir(input_path):
print(f"错误：路径「{input_path}」不是有效文件夹")
return []
# 创建输出文件夹（存放在输入路径下，命名为class_diagrams）
output_dir = os.path.join(input_path, "class_diagrams")
os.makedirs(output_dir, exist_ok=True)
# 筛选所有.py文件（排除__init__.py等特殊文件）
py_files = [
os.path.join(root, f)
for root, _, files in os.walk(input_path)
for f in files if f.endswith('.py') and not f.startswith('__')
]
# 批量生成类图
generated_diagrams = []
for py_file in py_files:
class_info, global_funs = extract_classes_and_functions_from_file(py_file)
# 若没有类，创建"模块函数"虚拟类存放全局函数
if not class_info:
class_info = {"模块函数": global_funs}
global_funs = []
# 生成输出文件名（原文件名+_class_diagram.png）
file_name = os.path.splitext(os.path.basename(py_file))[0]
output_file = os.path.join(output_dir, f"{file_name}_class_diagram.png")
# 绘制并保存类图
plot_class_diagram(class_info, global_funs, output_file, os.path.basename(py_file))
generated_diagrams.append(output_file)
print(f"\n所有类图已保存至：{output_dir}")
return generated_diagrams
if __name__ == "__main__":
# 交互输入文件夹路径，生成类图
input_folder = input("请输入Python文件所在文件夹路径：")
result = generate_class_diagram(input_folder)
# 打印生成结果
if result:
print("\n生成成功的类图列表：")
for idx, path in enumerate(result, 1):
print(f"{idx}. {path}")
else:
print("\n未生成任何类图")