Module `opshin.rewrite.rewrite_remove_type_stuff`

Expand source code

from typing import Optional

from ..typed_ast import (
    TypedAssign,
    ClassType,
    InstanceType,
    PolymorphicFunctionType,
    TypeInferenceError,
)
from ..util import CompilingNodeTransformer

"""
Remove class reassignments without constructors and polymorphic function reassignments

Both of these are only present during the type inference and are discarded or generated in-place during compilation.
"""


class RewriteRemoveTypeStuff(CompilingNodeTransformer):
    step = "Removing class and polymorphic function re-assignments"

    def visit_Assign(self, node: TypedAssign) -> Optional[TypedAssign]:
        assert (
            len(node.targets) == 1
        ), "Assignments to more than one variable not supported yet"
        try:
            if isinstance(node.value.typ, ClassType):
                try:
                    typ = node.value.typ.constr_type()
                except TypeInferenceError:
                    # no constr_type is also fine
                    return None
            else:
                typ = node.value.typ
            if isinstance(typ, InstanceType) and isinstance(
                typ.typ, PolymorphicFunctionType
            ):
                return None
        except AttributeError:
            # untyped attributes are fine too
            pass
        return node

Classes

class RewriteRemoveTypeStuff

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 RewriteRemoveTypeStuff(CompilingNodeTransformer):
    step = "Removing class and polymorphic function re-assignments"

    def visit_Assign(self, node: TypedAssign) -> Optional[TypedAssign]:
        assert (
            len(node.targets) == 1
        ), "Assignments to more than one variable not supported yet"
        try:
            if isinstance(node.value.typ, ClassType):
                try:
                    typ = node.value.typ.constr_type()
                except TypeInferenceError:
                    # no constr_type is also fine
                    return None
            else:
                typ = node.value.typ
            if isinstance(typ, InstanceType) and isinstance(
                typ.typ, PolymorphicFunctionType
            ):
                return None
        except AttributeError:
            # untyped attributes are fine too
            pass
        return node

Ancestors

CompilingNodeTransformer
TypedNodeTransformer
ast.NodeTransformer
ast.NodeVisitor

Class variables

var step

Methods

def visit(self, node)

Inherited from: CompilingNodeTransformer.visit

Visit a node.

def visit_Assign(self, node: TypedAssign) ‑> Optional[TypedAssign]

Expand source code

def visit_Assign(self, node: TypedAssign) -> Optional[TypedAssign]:
    assert (
        len(node.targets) == 1
    ), "Assignments to more than one variable not supported yet"
    try:
        if isinstance(node.value.typ, ClassType):
            try:
                typ = node.value.typ.constr_type()
            except TypeInferenceError:
                # no constr_type is also fine
                return None
        else:
            typ = node.value.typ
        if isinstance(typ, InstanceType) and isinstance(
            typ.typ, PolymorphicFunctionType
        ):
            return None
    except AttributeError:
        # untyped attributes are fine too
        pass
    return node