Как я могу заменить OrderedDict на dict в Python AST перед буквальным eval?

У меня есть строка с кодом Python, которую я мог бы оценить как Python с literal_eval, если бы в ней были только экземпляры OrderedDict, замененные на {}.

Я пытаюсь использовать ast.parse и ast.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)
3 6
python abstract-syntax-tree

3 ответа:

Идея состоит в том, чтобы заменить все OrderedDict узлы, представленные как ast.Call, имеющие определенные атрибуты (которые можно увидеть из ordered_dict_conditions ниже), на ast.Dict узлы, у которых 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.Tuple s (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).

2

Вы могли бы использовать ast.NodeVisitor класс для наблюдения за деревом OrderedDict, чтобы построить дерево {} вручную из найденных узлов, используя разбираемые узлы из пустого dict в качестве основы.

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 в качестве значений. Однако расширение более сложных структур должно быть возможно аналогичным образом.
0

Вместо того, чтобы использовать ast для разбора и преобразования выражения, вы можете также использовать регулярное выражение для этого. Например: 

>>> re.sub(
...     r"OrderedDict\(\[((\(('[a-z]+'), (\d+)\)),?\s*)+\]\)",
...     r'{\3: \4}',
...     "[OrderedDict([('a', 1)])]"
... )
"[{'a': 1}]"
Приведенное выше выражение основано на примере строки OP и рассматривает одиночные строки в кавычках как ключи и положительные целые числа как значения, но, конечно, оно может быть распространено на более сложные случаи.
0