Logo Questions Linux Laravel Mysql Ubuntu Git Menu
 

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

I have a string with Python code in it that I could evaluate as Python with literal_eval if it only had instances of OrderedDict replaced with {}.

I am trying to use ast.parse and ast.NodeTransformer to do the replacement, but when I catch the node with nodetype == 'Name' and node.id == 'OrderedDict', I can't find the list that is the argument in the node object so that I can replace it with a Dict node.

Is this even the right approach?

Some code:

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)
like image 318
Jim Hunziker Avatar asked Jul 09 '18 19:07

Jim Hunziker


1 Answers

The idea is to replace all OrderedDict nodes, represented as ast.Call having specific attributes (which can be seen from ordered_dict_conditions below), with ast.Dict nodes whose key / value arguments are extracted from the ast.Call arguments.

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}}

Notes

Because we want to replace the innermost node first, we place a call to super() at the beginning of the function.

Whenever an OrderedDict node is encountered, the following things are used:

  • node.args is a list containing the arguments to the OrderedDict(...) call.
  • This call has a single argument, namely a list containing key-value pairs as tuples, which is accessible by node.args[0] (ast.List) and node.args[0].elts are the tuples wrapped in a list.
  • So node.args[0].elts[i] are the different ast.Tuples (for i in range(len(node.args[0].elts))) whose elements are accessible again via the .elts attribute.
  • Finally node.args[0].elts[i].elts[0] are the keys and node.args[0].elts[i].elts[1] are the values which are used in the OrderedDict call.

The latter keys and values are then used to create a fresh ast.Dict instance which is then used to replace the current node (which was ast.Call).

like image 73
Leo K Avatar answered Oct 20 '22 18:10

Leo K