如何在 literal_eval 之前的 Python AST 中用 dict 替换 OrderedDict?

How can I replace OrderedDict with dict in a Python AST before literal_eval?

我有一个带有 Python 代码的字符串,如果它只有 OrderedDict 的实例替换为 [=15],我可以将其评估为 Python 和 literal_eval =].

我正在尝试使用 ast.parseast.NodeTransformer 进行替换,但是当我使用 nodetype == 'Name' and node.id == 'OrderedDict' 捕获节点时,我找不到作为参数的列表在节点对象中,以便我可以用 Dict 节点替换它。

这是正确的方法吗?

一些代码:

from ast import NodeTransformer, parse

py_str = "[OrderedDict([('a', 1)])]"

class Transformer(NodeTransformer):
    def generic_visit(self, node):
        nodetype = type(node).__name__

        if nodetype == 'Name' and node.id == 'OrderedDict':
            pass # ???

        return NodeTransformer.generic_visit(self, node)

t = Transformer()

tree = parse(py_str)

t.visit(tree)

您可以使用 ast.NodeVisitor class 观察 OrderedDict 树,以便使用从中解析的节点从遇到的节点手动构建 {} 树空字典作为基础。

import ast
from collections import deque


class Builder(ast.NodeVisitor):
    def __init__(self):
        super().__init__()
        self._tree = ast.parse('[{}]')
        self._list_node = self._tree.body[0].value
        self._dict_node = self._list_node.elts[0]
        self._new_item = False

    def visit_Tuple(self, node):
        self._new_item = True
        self.generic_visit(node)

    def visit_Str(self, node):
        if self._new_item:
            self._dict_node.keys.append(node)
        self.generic_visit(node)

    def visit_Num(self, node):
        if self._new_item:
            self._dict_node.values.append(node)
            self._new_item = False
        self.generic_visit(node)

    def literal_eval(self):
        return ast.literal_eval(self._list_node)


builder = Builder()
builder.visit(ast.parse("[OrderedDict([('a', 1)])]"))
print(builder.literal_eval())

请注意,这仅适用于使用 str 作为键并使用 int 作为值的示例的简单结构。然而,以类似的方式扩展到更复杂的结构应该是可能的。

想法是用 ast.Dict 节点替换所有 OrderedDict 节点,表示为具有特定属性(可以从下面的 ordered_dict_conditions 看出)的 ast.Call key / value 个参数是从 ast.Call 个参数中提取的。

import ast


class Transformer(ast.NodeTransformer):
    def generic_visit(self, node):
        # Need to call super() in any case to visit child nodes of the current one.
        super().generic_visit(node)
        ordered_dict_conditions = (
            isinstance(node, ast.Call)
            and isinstance(node.func, ast.Name)
            and node.func.id == 'OrderedDict'
            and len(node.args) == 1
            and isinstance(node.args[0], ast.List)
        )
        if ordered_dict_conditions:
            return ast.Dict(
                [x.elts[0] for x in node.args[0].elts],
                [x.elts[1] for x in node.args[0].elts]
            )
        return node


def transform_eval(py_str):
    return ast.literal_eval(Transformer().visit(ast.parse(py_str, mode='eval')).body)


print(transform_eval("[OrderedDict([('a', 1)]), {'k': 'v'}]"))  # [{'a': 1}, {'k': 'v'}]
print(transform_eval("OrderedDict([('a', OrderedDict([('b', 1)]))])"))  # {'a': {'b': 1}}

备注

因为我们想先替换最里面的节点,所以我们在函数的开头调用了super()

每当遇到 OrderedDict 节点时,将使用以下内容:

  • node.args 是一个列表,其中包含 OrderedDict(...) 调用的参数。
  • 这个调用有一个参数,即包含键值对作为元组的列表,可以通过 node.args[0] (ast.List) 访问,node.args[0].elts 是包裹在中的元组list.
  • 因此 node.args[0].elts[i] 是不同的 ast.Tuples (for i in range(len(node.args[0].elts))),其元素可通过 .elts 属性再次访问。
  • 最后 node.args[0].elts[i].elts[0] 是键,node.args[0].elts[i].elts[1]OrderedDict 调用中使用的值。

后面的键和值随后用于创建一个新的 ast.Dict 实例,然后用于替换当前节点(ast.Call)。

除了使用 ast 来解析和转换表达式之外,您还可以使用正则表达式来执行此操作。例如:

>>> re.sub(
...     r"OrderedDict\(\[((\(('[a-z]+'), (\d+)\)),?\s*)+\]\)",
...     r'{: }',
...     "[OrderedDict([('a', 1)])]"
... )
"[{'a': 1}]"

上述表达式基于OP的示例字符串,将单引号字符串视为键,将正整数视为值,但当然可以扩展到更复杂的情况。