用 LibCST 编写 codemod 替换 datetime.utcnow()

问题场景

项目里散落着 datetime.utcnow()。它返回的是不带时区信息的 naive datetime，和 datetime.now(timezone.utc) 返回的 aware datetime 混用后，容易出现比较、序列化和跨时区展示不一致的问题。手工替换风险在于：文件多、容易漏改，补 from datetime import timezone 时又可能破坏文件头注释、docstring 和 __future__ 导入顺序。

目标很明确：把所有 datetime.utcnow() 改为 datetime.now(timezone.utc)，在需要时自动补充 from datetime import timezone，并尽量保留原有格式与注释。

为什么选 LibCST

正则替换很难可靠处理注释、括号空白和 import 位置。LibCST 把源码解析成保留格式信息的 CST，转换后再生成代码，适合做这类小范围 codemod。下面的脚本只处理一种明确写法，避免过度改写。

最小脚本

保存为 tools/codemod_utcnow.py，先安装依赖：

pip install libcst

脚本内容：

#!/usr/bin/env python3
from __future__ import annotations

import argparse
import sys
from pathlib import Path

import libcst as cst


def _attr(value: cst.BaseExpression, name: str) -> cst.Attribute:
    return cst.Attribute(value=value, attr=cst.Name(name))


class RewriteUtcnow(cst.CSTTransformer):
    def __init__(self) -> None:
        self.needs_timezone = False

    def leave_Call(self, original_node: cst.Call, updated_node: cst.Call) -> cst.Call:
        func = original_node.func
        if not (
            isinstance(func, cst.Attribute)
            and isinstance(func.value, cst.Name)
            and func.value.value == 'datetime'
            and isinstance(func.attr, cst.Name)
            and func.attr.value == 'utcnow'
        ):
            return updated_node
        if original_node.args or original_node.keywords:
            return updated_node
        self.needs_timezone = True
        return updated_node.with_changes(
            func=_attr(cst.Name('datetime'), 'now'),
            args=[cst.Arg(_attr(cst.Name('timezone'), 'utc'))],
        )

    def leave_Module(self, original_node: cst.Module, updated_node: cst.Module) -> cst.Module:
        if not self.needs_timezone or _has_top_level_timezone(updated_node):
            return updated_node
        body = list(updated_node.body)
        stmt = cst.parse_statement('from datetime import timezone\n')
        body.insert(_import_index(body), stmt)
        return updated_node.with_changes(body=body)


def _has_top_level_timezone(module: cst.Module) -> bool:
    for stmt in module.body:
        if not isinstance(stmt, cst.SimpleStatementLine):
            continue
        for small in stmt.body:
            if not isinstance(small, cst.ImportFrom):
                continue
            if not (isinstance(small.module, cst.Name) and small.module.value == 'datetime'):
                continue
            for alias in small.names:
                if alias.name.value != 'timezone':
                    continue
                if alias.asname is None or alias.asname.name.value == 'timezone':
                    return True
    return False


def _is_docstring(stmt: cst.BaseStatement) -> bool:
    return (
        isinstance(stmt, cst.SimpleStatementLine)
        and len(stmt.body) == 1
        and isinstance(stmt.body[0], cst.Expr)
        and isinstance(stmt.body[0].value, cst.SimpleString)
    )


def _is_future_import(stmt: cst.BaseStatement) -> bool:
    if not isinstance(stmt, cst.SimpleStatementLine) or not stmt.body:
        return False
    return all(
        isinstance(s, cst.ImportFrom)
        and isinstance(s.module, cst.Name)
        and s.module.value == '__future__'
        for s in stmt.body
    )


def _import_index(body: list[cst.BaseStatement]) -> int:
    i = 0
    if body and _is_docstring(body[0]):
        i = 1
    while i < len(body) and _is_future_import(body[i]):
        i += 1
    return i


def transform(source: str) -> tuple[str, bool]:
    module = cst.parse_module(source)
    new_module = module.visit(RewriteUtcnow())
    out = new_module.code
    return out, out != source


def iter_py(paths: list[str]):
    for raw in paths:
        p = Path(raw)
        if p.is_file() and p.suffix == '.py':
            yield p
        elif p.is_dir():
            yield from sorted(p.rglob('*.py'))


def main(argv: list[str] | None = None) -> int:
    parser = argparse.ArgumentParser()
    parser.add_argument('paths', nargs='+')
    parser.add_argument('--check', action='store_true', help='只检查，不写文件')
    parser.add_argument('--dry-run', action='store_true', help='只打印会修改的文件')
    args = parser.parse_args(argv)

    changed_files = []
    parse_errors = []
    for path in iter_py(args.paths):
        source = path.read_text(encoding='utf-8')
        try:
            output, changed = transform(source)
        except cst.ParserSyntaxError as exc:
            parse_errors.append(path)
            print(f'parse error: {path}: {exc}', file=sys.stderr)
            continue
        if not changed:
            continue
        changed_files.append(path)
        if args.check:
            print(f'need rewrite: {path}')
        elif args.dry_run:
            print(f'would rewrite: {path}')
        else:
            path.write_text(output, encoding='utf-8')
            print(f'rewritten: {path}')

    if args.check and changed_files:
        print(f'{len(changed_files)} file(s) need rewrite', file=sys.stderr)
        return 1
    if parse_errors:
        return 2
    return 0


if __name__ == '__main__':
    raise SystemExit(main())

脚本用 leave_Call 定位调用，用 with_changes 保留调用节点周围的大部分空白；leave_Module 只在确实发生过替换且顶层没有可用 timezone 导入时插入新行。导入插入位置会跳过模块 docstring 和 __future__ 导入，减少对文件头结构的破坏。

运行与 dry-run

先看会改哪些文件，不落盘：

python tools/codemod_utcnow.py --dry-run src tests

确认后执行：

python tools/codemod_utcnow.py src tests

建议在独立分支运行，执行后查看 diff，重点检查文件头、docstring 和 import 区域。

在 CI 中用 --check 拦截遗漏

--check 不写文件，只要发现某个文件会被改写就以退出码 1 结束，适合放在 CI 或流水线里：

- run: pip install libcst
- run: python tools/codemod_utcnow.py --check src tests

这样新提交如果又写回 datetime.utcnow()，或者缺少 timezone 导入，流水线会直接失败。也可以把同一命令配到 pre-commit 中，在本地提前拦截。

边界与限制

• 只匹配 datetime.utcnow() 且参数为空的调用；datetime.datetime.utcnow()、dt.utcnow() 这类别名写法不会命中。
• 基于语法匹配，不解析名称绑定；如果 datetime 被局部变量遮蔽，脚本无法识别。
• 空括号内夹注释等极端写法可能无法完全保留，普通行尾注释和缩进会保留。
• 脚本按 UTF-8 读写；非 UTF-8 文件需先转换或跳过。
• 插入的是顶层 from datetime import timezone；若项目规范要求 import 排序，可在 codemod 后统一跑格式化工具。

把范围限定清楚、用 --check 固化为检查项，比一次性大规模正则替换更可控。