我用Python让AI自动给代码写文档，终于不用手动补了

# pip install openai

import ast
import os
import textwrap
from openai import OpenAI

# 建连接
client = OpenAI(api_key="your_api_key_here")


# ==================================================
# 第一步：用ast解析Python源码，提取函数和类信息
# ==================================================

def extract_functions(source_code):
    """从源码里把所有函数定义扒出来"""
    tree = ast.parse(source_code)
    functions = []

    for node in ast.walk(tree):
        if not isinstance(node, (ast.FunctionDef, ast.AsyncFunctionDef)):
            continue

        # 跳过私有函数（以_开头的）
        if node.name.startswith("_")  and not node.name.startswith("__"):
            continue

        # 提取参数信息
        args_info = []
        for arg in node.args.args:
            arg_name = arg.arg
            # 如果有类型注解，也拿出来
            arg_type = ""
            if arg.annotation:
                arg_type = ast.unparse(arg.annotation)
            args_info.append({"name": arg_name, "type": arg_type})

        # 提取返回值注解
        return_type = ""
        if node.returns:
            return_type = ast.unparse(node.returns)

        # 看看函数体有没有docstring
        existing_doc = ast.get_docstring(node)

        # 把函数源码也提取出来，给大模型参考
        func_source = ast.unparse(node)

        functions.append({
            "name": node.name,
            "args": args_info,
            "return_type": return_type,
            "existing_doc": existing_doc,
            "source": func_source,
            "lineno": node.lineno,
        })

    return functions


def extract_classes(source_code):
    """提取类的信息"""
    tree = ast.parse(source_code)
    classes = []

    for node in ast.walk(tree):
        if not isinstance(node, ast.ClassDef):
            continue

        methods = []
        for item in node.body:
            if isinstance(item, (ast.FunctionDef, ast.AsyncFunctionDef)):
                methods.append(item.name)

        classes.append({
            "name": node.name,
            "methods": methods,
            "bases": [ast.unparse(b) forbinnode.bases],
            "existing_doc": ast.get_docstring(node),
        })

    return classes


# ==================================================
# 第二步：让大模型给每个函数生成docstring
# ==================================================

def generate_docstring(func_info):
    """给一个函数生成docstring"""

    # 如果已经有docstring了，可以跳过
    if func_info["existing_doc"]:
        return func_info["existing_doc"]

    args_desc = ""
    for arg in func_info["args"]:
        if arg["name"] == "self":
            continue
        type_str = f" ({arg['type']})"ifarg["type"] else""
        args_desc += f"  - {arg['name']}{type_str}: \n"

    prompt = f"""请给下面这个Python函数写一个docstring。

函数源码：
```python
{func_info['source']}

prompt = f"""请给下面这个Python类写一段文档说明。

filename = os.path.basename(file_path)
md = f"# {filename} 接口文档\n\n"
md += f"> 自动生成，源码路径: `{file_path}`\n\n"

# 类文档
if classes:
    md += "## 类\n\n"
    for cls in classes:
        md += f"### {cls['name']}\n\n"
        doc = generate_class_docstring(cls)
        md += f"{doc}\n\n"
        md += f"包含方法: {', '.join(cls['methods'])}\n\n"

# 函数文档
if functions:
    md += "## 函数\n\n"
    for func, doc in zip(functions, docstrings):
        md += f"### `{func['name']}`\n\n"
        md += f"```python\n{func['source'][:200]}...\n```\n\n"
        md += f"{doc}\n\n"

        # 参数列表
        if func["args"]:
            md += "**参数:**\n"
            for arg in func["args"]:
                if arg["name"] == "self":
                    continue
                type_str = f" (`{arg['type']}`)" if arg["type"] else ""
                md += f"- `{arg['name']}`{type_str}\n"
            md += "\n"

        if func["return_type"]:
            md += f"**返回类型:** `{func['return_type']}`\n\n"

        md += "---\n\n"

return md

print(f"正在处理: {file_path}")

with open(file_path, "r", encoding="utf-8") as f:
    source = f.read()

# 提取函数和类
functions = extract_functions(source)
classes = extract_classes(source)

print(f"  找到 {len(functions)} 个函数, {len(classes)} 个类")

# 给每个函数生成docstring
docstrings = []
for i, func in enumerate(functions):
    print(f"  正在生成 [{i+1}/{len(functions)}] {func['name']}...")
    doc = generate_docstring(func)
    docstrings.append(doc)

# 生成Markdown文档
md_content = generate_markdown(file_path, functions, classes, docstrings)

# 保存
output_path = file_path.replace(".py", "_docs.md")
with open(output_path, "w", encoding="utf-8") as f:
    f.write(md_content)

print(f"  文档已保存到: {output_path}")
return output_path

py_files = []
for root, dirs, files in os.walk(folder_path):
    # 跳过隐藏目录和常见的非源码目录
    dirs[:] = [d for d in dirs if not d.startswith((".", "__"))]
    for f in files:
        if f.endswith(".py") and not f.startswith("test_"):
            py_files.append(os.path.join(root, f))

print(f"共找到 {len(py_files)} 个Python文件\n")

for fp in py_files:
    try:
        document_python_file(fp)
    except Exception as e:
        print(f"  处理失败: {e}")
    print()

# 或者批量处理整个项目
# batch_document("./my_project")

## 运行效果

我拿项目里一个数据处理的文件试了一下，里面有12个函数、2个类。

生成的文档长这样（截取一部分）：

```markdown
# data_processor.py 接口文档

> 自动生成，源码路径: `data_processor.py`

## 函数

### `load_csv_data`

加载CSV文件并做基础清洗。自动处理编码问题，优先用utf-8，失败了换gbk。

Args:
  filepath (str): CSV文件路径
  encoding (str): 文件编码，默认utf-8

Returns:
  pd.DataFrame: 清洗后的数据

Raises:
  FileNotFoundError: 文件不存在时抛出