Module opshin.rewrite.rewrite_tuple_assign

Classes

class RewriteTupleAssign

A :class:NodeVisitor subclass that walks the abstract syntax tree and allows modification of nodes.

The NodeTransformer will walk the AST and use the return value of the visitor methods to replace or remove the old node. If the return value of the visitor method is None, the node will be removed from its location, otherwise it is replaced with the return value. The return value may be the original node in which case no replacement takes place.

Here is an example transformer that rewrites all occurrences of name lookups (foo) to data['foo']::

class RewriteName(NodeTransformer):

   def visit_Name(self, node):
       return Subscript(
           value=Name(id='data', ctx=Load()),
           slice=Constant(value=node.id),
           ctx=node.ctx
       )

Keep in mind that if the node you're operating on has child nodes you must either transform the child nodes yourself or call the :meth:generic_visit method for the node first.

For nodes that were part of a collection of statements (that applies to all statement nodes), the visitor may also return a list of nodes rather than just a single node.

Usually you use the transformer like this::

node = YourTransformer().visit(node)

Expand source code

class RewriteTupleAssign(CompilingNodeTransformer):
    step = "Rewriting tuple deconstruction in assignments"

    unique_id = 0

    def visit_Assign(self, node: Assign) -> typing.List[stmt]:
        if not isinstance(node.targets[0], Tuple):
            return [node]
        uid = self.unique_id
        self.unique_id += 1
        assignments = [Assign([Name(f"2_{uid}_tup", Store())], self.visit(node.value))]
        for i, t in enumerate(node.targets[0].elts):
            assignments.append(
                Assign(
                    [t],
                    Subscript(
                        value=Name(f"2_{uid}_tup", Load()),
                        slice=Constant(i),
                        ctx=Load(),
                    ),
                )
            )
        # recursively resolve multiple layers of tuples
        transformed = sum([self.visit(a) for a in assignments], [])
        return transformed

    def visit_For(self, node: For) -> For:
        # rewrite deconstruction in for loops
        if not isinstance(node.target, Tuple):
            return self.generic_visit(node)
        new_for = copy(node)
        new_for.iter = self.visit(node.iter)
        uid = self.unique_id
        self.unique_id += 1
        # write the tuple into a singleton variable
        new_for.target = Name(f"2_{uid}_tup", Store())
        assignments = []
        # iteratively assign the deconstructed parts to the original variable names
        for i, t in enumerate(node.target.elts):
            assignments.append(
                Assign(
                    [t],
                    Subscript(
                        value=Name(f"2_{uid}_tup", Load()),
                        slice=Constant(i),
                        ctx=Load(),
                    ),
                )
            )
        new_for.body = assignments + node.body
        # recursively resolve multiple layers of tuples
        # further layers should be handled by the normal tuple assignment though
        return self.visit(new_for)

Ancestors

• CompilingNodeTransformer
• TypedNodeTransformer
• ast.NodeTransformer
• ast.NodeVisitor

Class variables

var step

var unique_id

Methods

def visit(self, node)

Inherited from: CompilingNodeTransformer.visit

Visit a node.

def visit_Assign(self, node: ast.Assign) ‑> List[ast.stmt]

Expand source code

def visit_Assign(self, node: Assign) -> typing.List[stmt]:
    if not isinstance(node.targets[0], Tuple):
        return [node]
    uid = self.unique_id
    self.unique_id += 1
    assignments = [Assign([Name(f"2_{uid}_tup", Store())], self.visit(node.value))]
    for i, t in enumerate(node.targets[0].elts):
        assignments.append(
            Assign(
                [t],
                Subscript(
                    value=Name(f"2_{uid}_tup", Load()),
                    slice=Constant(i),
                    ctx=Load(),
                ),
            )
        )
    # recursively resolve multiple layers of tuples
    transformed = sum([self.visit(a) for a in assignments], [])
    return transformed

def visit_For(self, node: ast.For) ‑> ast.For

Expand source code

def visit_For(self, node: For) -> For:
    # rewrite deconstruction in for loops
    if not isinstance(node.target, Tuple):
        return self.generic_visit(node)
    new_for = copy(node)
    new_for.iter = self.visit(node.iter)
    uid = self.unique_id
    self.unique_id += 1
    # write the tuple into a singleton variable
    new_for.target = Name(f"2_{uid}_tup", Store())
    assignments = []
    # iteratively assign the deconstructed parts to the original variable names
    for i, t in enumerate(node.target.elts):
        assignments.append(
            Assign(
                [t],
                Subscript(
                    value=Name(f"2_{uid}_tup", Load()),
                    slice=Constant(i),
                    ctx=Load(),
                ),
            )
        )
    new_for.body = assignments + node.body
    # recursively resolve multiple layers of tuples
    # further layers should be handled by the normal tuple assignment though
    return self.visit(new_for)