Module opshin.rewrite.rewrite_import_uplc_builtins
Expand source code
import re
from copy import copy
from typing import Optional
from enum import Enum
from ..util import CompilingNodeTransformer
from ..typed_ast import *
"""
Checks that there was an import of wraps_builtin if there are any wrapped builtins
"""
DECORATOR_NAME = "wraps_builtin"
class RewriteImportUPLCBuiltins(CompilingNodeTransformer):
step = "Resolving imports and usage of UPLC builtins"
imports_uplc_builtins = False
def visit_ImportFrom(self, node: ImportFrom) -> Optional[AST]:
if node.module != "opshin.bridge":
return node
for n in node.names:
assert (
n.name == DECORATOR_NAME
), "Imports something other from the bridge than the builtin wrapper"
assert n.asname is None, "Renames the builtin wrapper. This is forbidden."
self.imports_uplc_builtins = True
return None
def visit_FunctionDef(self, node: TypedFunctionDef) -> AST:
if not node.decorator_list or len(node.decorator_list) != 1:
return node
is_wrapped = any(isinstance(n, Name) and n.id for n in node.decorator_list)
if not is_wrapped:
return node
assert (
self.imports_uplc_builtins
), "To wrap builtin functions, you need to import the builtin function. Add `from opshin.bridge import wraps_builtin` to your code."
# we replace the body with a forwarded call to the wrapped builtin
CamelCaseFunName = "".join(
p.capitalize() for p in re.split(r"_(?!\d)", node.orig_name)
)
uplc_fun = plt.__dict__[CamelCaseFunName]
pluto_expression = RawPlutoExpr(
typ=node.typ.typ.rettyp,
expr=uplc_fun(
*(plt.Force(plt.Var(a.arg)) for a in node.args.args),
),
)
node_cp = copy(node)
node_cp.body = [Return(pluto_expression, typ=node.typ.typ.rettyp)]
return node_cp
Classes
class RewriteImportUPLCBuiltins
-
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 isNone
, 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
) todata['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 RewriteImportUPLCBuiltins(CompilingNodeTransformer): step = "Resolving imports and usage of UPLC builtins" imports_uplc_builtins = False def visit_ImportFrom(self, node: ImportFrom) -> Optional[AST]: if node.module != "opshin.bridge": return node for n in node.names: assert ( n.name == DECORATOR_NAME ), "Imports something other from the bridge than the builtin wrapper" assert n.asname is None, "Renames the builtin wrapper. This is forbidden." self.imports_uplc_builtins = True return None def visit_FunctionDef(self, node: TypedFunctionDef) -> AST: if not node.decorator_list or len(node.decorator_list) != 1: return node is_wrapped = any(isinstance(n, Name) and n.id for n in node.decorator_list) if not is_wrapped: return node assert ( self.imports_uplc_builtins ), "To wrap builtin functions, you need to import the builtin function. Add `from opshin.bridge import wraps_builtin` to your code." # we replace the body with a forwarded call to the wrapped builtin CamelCaseFunName = "".join( p.capitalize() for p in re.split(r"_(?!\d)", node.orig_name) ) uplc_fun = plt.__dict__[CamelCaseFunName] pluto_expression = RawPlutoExpr( typ=node.typ.typ.rettyp, expr=uplc_fun( *(plt.Force(plt.Var(a.arg)) for a in node.args.args), ), ) node_cp = copy(node) node_cp.body = [Return(pluto_expression, typ=node.typ.typ.rettyp)] return node_cp
Ancestors
- CompilingNodeTransformer
- TypedNodeTransformer
- ast.NodeTransformer
- ast.NodeVisitor
Class variables
var imports_uplc_builtins
var step
Methods
def visit(self, node)
-
Inherited from:
CompilingNodeTransformer
.visit
Visit a node.
def visit_FunctionDef(self, node: TypedFunctionDef) ‑> ast.AST
def visit_ImportFrom(self, node: ast.ImportFrom) ‑> ast.AST | None