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

上个月公司来了个新同事，叫小陈，刚毕业一年。leader让他接手我去年写的一个数据处理项目，里面有60多个Python文件，200多个函数。

入职第三天，小陈来找我。“王哥，这个process_data_v3函数是干嘛的？参数flag是什么意思？”

我看了一眼那个函数。说实话，我也忘了。那是去年赶工写的，注释就一行# 处理数据，参数名也随便起的。我自己都看不懂，更别说别人了。

leader知道这事以后，说了一句让我心里一凉的话：“趁这周不忙，把项目文档补一下吧。”

200多个函数。每个都要写docstring，说明参数、返回值、异常情况。再加上README、API文档。我粗算了一下，纯手动写，少说要三四天。

但转念一想，Python有ast模块，能把源码解析成结构化的语法树。大模型又擅长写文档。这俩凑一起，能不能让AI自动帮我生成文档？

花了一个下午，写了个自动文档生成器。核心思路就是用ast把每个函数和类的签名信息提取出来，发给大模型让它写docstring，然后回写到源码里。

完整代码

# pip install openai
importast
importos
importtextwrap
fromopenaiimportOpenAI
# 建连接
client = OpenAI(api_key="your_api_key_here")

# ==================================================
# 第一步：用ast解析Python源码，提取函数和类信息
# ==================================================
defextract_functions(source_code):
    """从源码里把所有函数定义扒出来"""
    tree = ast.parse(source_code)
    functions = []
    fornodeinast.walk(tree):
        ifnotisinstance(node, (ast.FunctionDef, ast.AsyncFunctionDef)):
            continue
        # 跳过私有函数（以_开头的）
        ifnode.name.startswith("_") andnotnode.name.startswith("__"):
            continue
        # 提取参数信息
        args_info = []
        forarginnode.args.args:
            arg_name = arg.arg
            # 如果有类型注解，也拿出来
            arg_type = ""
            ifarg.annotation:
                arg_type = ast.unparse(arg.annotation)
            args_info.append({"name": arg_name, "type": arg_type})
        # 提取返回值注解
        return_type = ""
        ifnode.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,
        })
    returnfunctions

defextract_classes(source_code):
    """提取类的信息"""
    tree = ast.parse(source_code)
    classes = []
    fornodeinast.walk(tree):
        ifnotisinstance(node, ast.ClassDef):
            continue
        methods = []
        foriteminnode.body:
            ifisinstance(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),
        })
    returnclasses

# ==================================================
# 第二步：让大模型给每个函数生成docstring
# ==================================================
defgenerate_docstring(func_info):
    """给一个函数生成docstring"""
    # 如果已经有docstring了，可以跳过
    iffunc_info["existing_doc"]:
        returnfunc_info["existing_doc"]
    args_desc = ""
    forarginfunc_info["args"]:
        ifarg["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']}

函数名：{func_info[’name’]}参数：{func_info[’args’]}返回类型：{func_info[’return_type’]}

要求：

 

用Google风格的docstring

 

 

用中文写

 

 

简洁明了，别写废话

 

 

包含Args、Returns、Raises部分（如果有的话）

 

 

只返回docstring内容，不要返回其他东西“””

 

resp = client.chat.completions.create(    model=“gpt-4o”,    messages=[{“role”: “user”, “content”: prompt}],    temperature=0.2  # 文档要准确，低温)

 

return resp.choices[0].message.content.strip()

 

def generate_class_docstring(class_info):“”“给类生成文档”“”

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

类名：{class_info[’name’]}父类：{class_info[’bases’]}包含的方法：{class_info[’methods’]}

要求：

 

用中文写

 

 

说清楚这个类是干什么用的

 

 

列出主要的公开方法

 

 

简洁，不超过150字“””

 

resp = client.chat.completions.create(    model=“gpt-4o”,    messages=[{“role”: “user”, “content”: prompt}],    temperature=0.2)

 

return resp.choices[0].message.content.strip()

 

========================================================

第三步：生成Markdown文档

========================================================

def generate_markdown(file_path, functions, classes, docstrings):“”“把文档整合成一个Markdown文件”“”

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

========================================================

主函数：串起来

========================================================

def document_python_file(file_path):“”“给一个Python文件生成完整文档”“”

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

def batch_document(folder_path):“”“批量处理一个文件夹下所有Python文件”“”

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()

if name == “main”:    # 单个文件    document_python_file(“data_processor.py”)

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

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

正在处理: data_processor.py  找到 12 个函数, 2 个类  正在生成 [1/12] load_csv_data…  正在生成 [2/12] clean_null_values…  正在生成 [3/12] merge_datasets…  正在生成 [4/12] calculate_statistics...…  文档已保存到: data_processor_docs.md

生成的文档长这样（截取一部分）：
```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: 文件不存在时抛出

12个函数，每个docstring都写得挺准。大模型看了函数的源码和参数信息，基本能猜出这个函数是干嘛的。比我去年随手写的# 处理数据强多了。

几个关键点

ast模块是核心。 Python自带的ast模块能把源码解析成语法树，你可以遍历这棵树找到所有的函数定义、类定义、参数信息。这比用正则匹配靠谱得多，因为正则处理不了嵌套和缩进。ast.unparse()还能把语法树节点还原成源码字符串，非常好用。

函数源码要一起发。 光告诉大模型函数名和参数名不够。比如process(data, flag)这种命名，不看源码根本不知道flag是什么意思。把整个函数体发过去，大模型就能从代码逻辑推断出参数含义。

跳过已有的docstring。 代码里有个判断，如果函数已经有docstring了就不重新生成。这样可以增量更新，不会覆盖你之前手写的文档。

批量处理跳过测试文件。batch_document函数会跳过test_开头的文件和隐藏目录，因为这些一般不需要生成文档。

后来的事

我把整个项目的文档都生成了一遍，200多个函数，大半天搞定。小陈拿着文档看了两天，大部分函数都能自己搞懂了，不用再来问我。

leader看效果不错，让我把这个工具加到CI流程里。每次提交代码自动更新文档，以后文档不会再欠债了。

不过有一点要注意。AI生成的文档偶尔会有理解偏差，特别是那种业务逻辑很复杂的函数。所以生成之后最好快速过一遍，花不了多少时间，总比从零手写快太多了。

以上就是“我用Python让AI自动给代码写文档，终于不用手动补了！”的详细内容，想要了解更多Python教程欢迎持续关注编程学习网。