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

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

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

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

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

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

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

完整代码

# pip install openaiimport astimport osimport textwrapfrom 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)# 看看函数体有没有docstringexisting_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 functionsdef extract_classes(source_code):"""提取类的信息"""tree = ast.parse(source_code)classes = []for node in ast.walk(tree):if not isinstance(node, ast.ClassDef):continuemethods = []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":continuetype_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’]}

要求：

1. 用Google风格的docstring

2. 用中文写

3. 简洁明了，别写废话

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

5. 只返回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’]}

要求：

1. 用中文写

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

3. 列出主要的公开方法

4. 简洁，不超过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)} 个类")# 给每个函数生成docstringdocstrings = []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-8Returns:  pd.DataFrame: 清洗后的数据Raises:  FileNotFoundError: 文件不存在时抛出

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

几个关键点

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

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

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

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

后来的事

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

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

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

https://ima.qq.com/wiki/?shareId=f2628818f0874da17b71ffa0e5e8408114e7dbad46f1745bbd1cc1365277631c