根据python代码自动生成类图的实现方法[附带python源码]

概述

利用python库抽象语法树(AST)和类图描述语言(PlantUML),实现自动将python代码生成类图的目的。

环境

windows+vscode+python+plantuml

✒️网上好像大部分都是用Pyreverse库来实现的，但是我实际测试发现只能在一个文件中才能行，当然应该有解决方法，由于我用plantuml居多，就想是不是能转换一下，所以这个方法就没多研究。此处也安利一把PlantUML，有兴趣的可以看下官网介绍，有比较详细的使用说明，

预览按钮在编辑界面的右上角

使用说明

这个脚本只是简单的处理继承关系和多个文件的自动化访问，还有很多方面的不足例如没有属性区分、错误处理、性能等。

在使用此脚本时需要将directory_to_process = '../hb'中的’…/hb’`更改为您需要查看类关系的代码目录，然后执行脚本后会生成对应的plantuml格式的文件，然后点击生成的plantuml的文件后点击右上角的预览，会出现一个处理中的界面，稍等后就可以看到对应的图片了。

python脚本源码

import ast
import os

class ClassInfo:
    def __init__(self, name, methods=None, attributes=None, bases=None):
        self.name = name
        self.methods = methods if methods is not None else []
        self.attributes = attributes if attributes is not None else []
        self.bases = bases if bases is not None else []

def parse_python_file(file_path):
    with open(file_path, 'r',encoding='utf-8') as file:
        source_code = file.read()
    
    tree = ast.parse(source_code)# 使用ast模块提供的parse()函数将源代码解析为抽象语法树对象
    classes = {}
    
    for node in ast.walk(tree): #使用walk()方法遍历该抽象语法树对象
        if isinstance(node, ast.ClassDef):
            class_name = node.name

            methods = [n.name for n in ast.walk(node) if isinstance(n, ast.FunctionDef)]
            attributes = []  # Note: Attributes are not directly represented in AST, usually set in __init__
             # 提取基类列表
            bases = []
            for base in node.bases:
                if isinstance(base, ast.Name):
                    # 如果基类是一个简单的名称，则提取其id
                    bases.append(base.id)
                else:
                    # 如果基类不是简单的名称，则可以选择跳过它、记录警告或错误
                    # 例如，记录警告：
                    # print(f"Warning: Non-simple base class expression found in class {class_name}: {ast.dump(base)}")
                    pass
            
            classes[class_name] = ClassInfo(class_name, methods, attributes, bases)
    
    return classes

def generate_plantuml(classes):
    plantuml_code = "@startuml\n"
    
    for class_name, info in classes.items():
        plantuml_code += f"class {class_name} {{\n"
        for method in info.methods:
            plantuml_code += f"  +{method}()\n"
        # Add attributes if you can extract them (not shown here)
        plantuml_code += "}\n"
        
        for base in info.bases:
            plantuml_code += f"{class_name} --|> {base}\n"
    
    plantuml_code += "@enduml\n"
    
    return plantuml_code

def process_directory(directory):
    all_classes = {}
    
    for root, dirs, files in os.walk(directory):
        for file in files:
            if file.endswith('.py'):
                file_path = os.path.join(root, file)
                classes = parse_python_file(file_path)
                all_classes.update(classes)
    
    return all_classes

# Example usage
directory_to_process = '../hb'
all_classes = process_directory(directory_to_process)
plantuml_desc = generate_plantuml(all_classes)

# Print or save the PlantUML description
print(plantuml_desc)
# You can save plantuml_desc to a .puml file and use PlantUML tools to convert it to an image.
with open('output.puml', 'w', encoding='utf-8') as file:
    file.write(plantuml_desc)