用 AI 写一个 Python 类型注解自动补全工具——让老代码享受类型安全

# typeai.py
import ast, click
from pathlib import Path
from openai import OpenAI
from dotenv import load_dotenv
load_dotenv()

client = OpenAI(api_key=os.getenv("DEEPSEEK_API_KEY"), base_url="https://api.deepseek.com/v1")

class TypeAnnotator(ast.NodeTransformer):
    """分析 AST，找出缺少类型注解的函数。"""

    def __init__(self):
        self.missing = []

    def visit_FunctionDef(self, node):
        # 检查函数参数是否有类型注解
        for arg in node.args.args:
            if arg.arg != "self" and arg.annotation is None:
                self.missing.append({
                    "file": "",
                    "func": node.name,
                    "line": node.lineno,
                    "code": ast.unparse(node).split(":")[0][:200],
                })
                break
        # 检查返回值是否有类型注解
        if node.returns is None:
            if not any(m["func"] == node.name for m in self.missing):
                self.missing.append({
                    "file": "",
                    "func": node.name,
                    "line": node.lineno,
                    "code": ast.unparse(node).split(":")[0][:200],
                })
        return self.generic_visit(node)

def find_missing_types(filepath: str) -> list:
    """扫描文件，找出缺少类型注解的函数。"""
    tree = ast.parse(Path(filepath).read_text())
    annotator = TypeAnnotator()
    annotator.visit(tree)
    for m in annotator.missing:
        m["file"] = filepath
    return annotator.missing

def generate_type_hints(code: str) -> str:
    """AI 生成类型注解。"""
    resp = client.chat.completions.create(
        model="deepseek-chat",
        messages=[{"role": "system", "content": "你是一个 Python 类型注解专家。为给定的函数签名添加准确的类型注解。只输出带注解的函数签名，不要解释。"},
                  {"role": "user", "content": f"添加类型注解：\n{code}"}],
        temperature=0.1,
        max_tokens=300,
    )
    return resp.choices[0].message.content

@click.command()
@click.argument("path", type=click.Path(exists=True))
@click.option("--apply", "-a", is_flag=True, help="自动应用修改")
@click.option("--dry-run", "-d", is_flag=True, help="只显示建议，不修改")
def main(path, apply, dry_run):
    """AI 自动补全 Python 类型注解。"""
    target = Path(path)
    files = [target] if target.is_file() else target.rglob("*.py")

    all_missing = []
    for f in files:
        if "test" in f.name or ".venv" in str(f):
            continue
        all_missing.extend(find_missing_types(str(f)))

    print(f"📊 发现 {len(all_missing)} 个缺少类型注解的函数")
    for m in all_missing[:20]:
        hint = generate_type_hints(m["code"])
        print(f"\n  📄 {m['file']}:{m['line']} {m['func']}")
        print(f"  当前：{m['code'][:100]}")
        print(f"  建议：{hint[:100]}")

if __name__ == "__main__":
    main()

# 扫描单个文件
python typeai.py src/api/auth.py --dry-run
# 📊 发现 3 个缺少类型注解的函数
#   📄 src/api/auth.py:12 login
#   当前：def login(email, password)
#   建议：def login(email: str, password: str) -> dict

# 扫描整个项目
python typeai.py src/ --dry-run