This package easy-ast
contains several utility about AST transformer.
pip install easy-ast
.
This example exchange all plus and mult.
import ast
from easy_ast import *
class PlusMultExchange(AstDecorator):
def visit_BinOp(self, node):
if isinstance(node.op, ast.Add):
result = ast.BinOp(left=node.left, op=ast.Mult(), right=node.right)
elif isinstance(node.op, ast.Mult):
result = ast.BinOp(left=node.left, op=ast.Add(), right=node.right)
else:
result = node
return self.generic_visit(result)
@PlusMultExchange()
def add(a, b):
return a + b
assert add(2, 3) == 6
Provide Statements
and Expression
to get the AST instance of code.
Exec
and Eval
is used to exec/eval code directly on AST
import ast
from easy_ast import *
@Expression
def tree():
(x + 1)**2
assert isinstance(tree, ast.Expression)
x = 6
assert Eval(tree) == 49
Because of the limit of python,
executing code will not effect variable in time, which may be solved in the future,
see pep558 and pep667.
Currently, the only way to use variable updated in exec is to visit it via locals()
manually,
or use it inside exec content directly.
import ast
from easy_ast import *
@Statements
def tree():
y = x + 1
z = y * 2
assert z == 6
assert isinstance(tree, ast.Module)
x = 2
Exec(tree)
assert locals()["z"] == 6
Class Macro
is used to create macro. Here is a simple example to do operator directly on python list.
import ast
from easy_ast import *
class List(Macro):
def __init__(self):
super().__init__()
self.symbol_current = 0
def visit_BinOp(self, node):
i = f"__list_loop_variable_{self.symbol_current}"
self.symbol_current += 1
j = f"__list_loop_variable_{self.symbol_current}"
self.symbol_current += 1
# Return [op(left, right) for i,j in zip(left, right)]
return ast.ListComp(
elt=ast.BinOp(
op=node.op,
left=ast.Name(id=i, ctx=ast.Load()),
right=ast.Name(id=j, ctx=ast.Load()),
),
generators=[
ast.comprehension(
target=ast.Tuple(elts=[
ast.Name(id=i, ctx=ast.Store()),
ast.Name(id=j, ctx=ast.Store()),
]),
iter=ast.Call(
func=ast.Name(id="zip", ctx=ast.Load()),
args=[
self.generic_visit(node.left),
self.generic_visit(node.right),
],
keywords=[],
),
ifs=[],
is_async=0,
)
],
)
a = [1, 2, 3]
b = [1, 2, 3]
@List().eval
def c():
a * b
assert c == [1, 4, 9]
This repository implements an AST transformer for Einstein notation based on Macro
for numpy array.
import numpy as np
from easy_ast.tensor_contract import TensorContract
b = np.array([1, 2])
c = np.array([1, 2])
expect = -np.array([[1, 2], [2, 4]])
@TensorContract().exec
def _(i, j):
a[i, j] = -b[i] * c[j]
assert np.all(a == expect)
a = np.random.randn(3, 2, 6, 5)
b = np.random.randn(3, 4)
c = np.random.randn(5, 4, 2)
d = np.einsum("ijk,il,klm->mj", a[:, 0], b, c)
assert d.shape == (2, 6)
r = np.zeros([6, 2, 2])
@TensorContract().exec
def _(i, j, k, l, m):
# i3, j6, k5, l4, m2
r[j, m, 1] = a[i, 0, j, k] * b[i, l] * c[k, l, m] - d[m, j]
assert np.sum(np.abs(r)) < 1e-6
It also supports non-standard Einstein notation.
import numpy as np
from easy_ast.tensor_contract import TensorContract
a = np.array([1, 2])
b = np.array([1, 2])
@TensorContract().exec
def _(i):
c[i] = a[i] * b[i]
assert np.all(c == [1, 4])
d = a[i]
assert d == 3