Python3.8环境下如何为ast.AST内置类动态添加属性方法？

以下两种方案均可实现需求，不需要修改内置类，无兼容性问题：

方案1：代理包装类（优先推荐） #

该方案完全规避内置类修改限制，兼容全Python版本，无副作用：

import ast
from functools import cached_property
from typing import Any

class ExtendedAST:
    def __init__(self, raw_node: ast.AST, parent: "ExtendedAST" = None):
        self._raw_node = raw_node
        self._parent = parent
        # 递归包装所有子节点
        for field, value in ast.iter_fields(raw_node):
            if isinstance(value, ast.AST):
                setattr(self, field, ExtendedAST(value, self))
            elif isinstance(value, list) and all(isinstance(item, ast.AST) for item in value):
                setattr(self, field, [ExtendedAST(item, self) for item in value])
            else:
                setattr(self, field, value)

    # 扩展属性：父节点
    @cached_property
    def parent(self):
        return self._parent

    # 扩展属性：直接子节点
    @cached_property
    def children(self):
        child_list = []
        for attr_value in self.__dict__.values():
            if isinstance(attr_value, ExtendedAST):
                child_list.append(attr_value)
            elif isinstance(attr_value, list):
                child_list.extend([item for item in attr_value if isinstance(item, ExtendedAST)])
        return child_list

    # 透传原生ast节点的所有属性、方法，保证API兼容
    def __getattr__(self, name: str) -> Any:
        return getattr(self._raw_node, name)

使用方式：将ast.parse返回的根节点传入包装类一次即可，所有子节点会自动完成包装：

raw_root = ast.parse("a = 1 + 2")
wrapped_root = ExtendedAST(raw_root)
# 正常调用原生属性
print(wrapped_root.body[0].targets[0].id) # 输出a
# 调用扩展属性
print(wrapped_root.body[0].parent is wrapped_root) # 输出True

方案2：弱引用字典存储扩展属性 #

如果不想引入包装层修改原有调用逻辑，可以使用该方案：

import ast
from weakref import WeakKeyDictionary
from functools import wraps, cached_property

# 弱引用字典存储节点扩展属性，不会干扰GC回收
_node_extra_store = WeakKeyDictionary()

def ast_extension(func):
    @wraps(func)
    def wrapper(node, *args, **kwargs):
        if node not in _node_extra_store:
            _node_extra_store[node] = {}
        return func(node, _node_extra_store[node], *args, **kwargs)
    # 注册到ast模块命名空间，全局可调用
    setattr(ast, func.__name__, wrapper)
    return wrapper

定义扩展属性/方法示例：

@ast_extension
def set_parent(node, ext_store, parent):
    ext_store["parent"] = parent

@ast_extension
def get_parent(node, ext_store):
    return ext_store.get("parent")

@ast_extension
@cached_property
def get_children(node, ext_store):
    return [n for n in ast.iter_child_nodes(node)]

使用方式：在首次遍历AST时为所有节点赋值父节点后，即可全局调用扩展方法：

root = ast.parse("a = 1 + 2")
# 首次遍历赋值父节点
for node in ast.walk(root):
    for child in ast.iter_child_nodes(node):
        ast.set_parent(child, node)
# 调用扩展方法
print(ast.get_parent(root.body[0]) is root) # 输出True

内容的提问来源于stack exchange，提问作者pawamoy