feat(gen): add simple generator
This commit is contained in:
144
midas/generator/generator.py
Normal file
144
midas/generator/generator.py
Normal file
@@ -0,0 +1,144 @@
|
|||||||
|
import ast
|
||||||
|
|
||||||
|
import midas.ast.python as p
|
||||||
|
|
||||||
|
|
||||||
|
class Generator(p.Stmt.Visitor[ast.stmt], p.Expr.Visitor[ast.expr]):
|
||||||
|
def __init__(self) -> None:
|
||||||
|
pass
|
||||||
|
|
||||||
|
def generate(self, stmts: list[p.Stmt]) -> str:
|
||||||
|
body: list[ast.stmt] = self._visit_body(stmts)
|
||||||
|
module = ast.Module(body=body, type_ignores=[])
|
||||||
|
module = ast.fix_missing_locations(module)
|
||||||
|
return ast.unparse(module)
|
||||||
|
|
||||||
|
def visit_binary_expr(self, expr: p.BinaryExpr) -> ast.expr:
|
||||||
|
return ast.BinOp(
|
||||||
|
left=expr.left.accept(self),
|
||||||
|
op=expr.operator,
|
||||||
|
right=expr.right.accept(self),
|
||||||
|
)
|
||||||
|
|
||||||
|
def visit_compare_expr(self, expr: p.CompareExpr) -> ast.expr:
|
||||||
|
return ast.Compare(
|
||||||
|
left=expr.left.accept(self),
|
||||||
|
ops=[expr.operator],
|
||||||
|
comparators=[expr.right.accept(self)],
|
||||||
|
)
|
||||||
|
|
||||||
|
def visit_unary_expr(self, expr: p.UnaryExpr) -> ast.expr:
|
||||||
|
return ast.UnaryOp(
|
||||||
|
op=expr.operator,
|
||||||
|
operand=expr.right.accept(self),
|
||||||
|
)
|
||||||
|
|
||||||
|
def visit_call_expr(self, expr: p.CallExpr) -> ast.expr:
|
||||||
|
return ast.Call(
|
||||||
|
func=expr.callee.accept(self),
|
||||||
|
args=[arg.accept(self) for arg in expr.arguments],
|
||||||
|
keywords=[
|
||||||
|
ast.keyword(arg=name, value=arg.accept(self))
|
||||||
|
for name, arg in expr.keywords.items()
|
||||||
|
],
|
||||||
|
)
|
||||||
|
|
||||||
|
def visit_get_expr(self, expr: p.GetExpr) -> ast.expr:
|
||||||
|
return ast.Attribute(
|
||||||
|
value=expr.object.accept(self),
|
||||||
|
attr=expr.name,
|
||||||
|
)
|
||||||
|
|
||||||
|
def visit_literal_expr(self, expr: p.LiteralExpr) -> ast.expr:
|
||||||
|
return ast.Constant(value=expr.value)
|
||||||
|
|
||||||
|
def visit_variable_expr(self, expr: p.VariableExpr) -> ast.expr:
|
||||||
|
return ast.Name(id=expr.name)
|
||||||
|
|
||||||
|
def visit_logical_expr(self, expr: p.LogicalExpr) -> ast.expr:
|
||||||
|
return ast.BoolOp(
|
||||||
|
op=expr.operator,
|
||||||
|
values=[expr.left.accept(self), expr.right.accept(self)],
|
||||||
|
)
|
||||||
|
|
||||||
|
def visit_cast_expr(self, expr: p.CastExpr) -> ast.expr:
|
||||||
|
# TODO: insert assertion
|
||||||
|
return expr.expr.accept(self)
|
||||||
|
|
||||||
|
def visit_ternary_expr(self, expr: p.TernaryExpr) -> ast.expr:
|
||||||
|
return ast.IfExp(
|
||||||
|
test=expr.test.accept(self),
|
||||||
|
body=expr.if_true.accept(self),
|
||||||
|
orelse=expr.if_false.accept(self),
|
||||||
|
)
|
||||||
|
|
||||||
|
def visit_list_expr(self, expr: p.ListExpr) -> ast.expr:
|
||||||
|
return ast.List(
|
||||||
|
elts=[item.accept(self) for item in expr.items],
|
||||||
|
)
|
||||||
|
|
||||||
|
def visit_subscript_expr(self, expr: p.SubscriptExpr) -> ast.expr:
|
||||||
|
return ast.Subscript(
|
||||||
|
value=expr.object.accept(self),
|
||||||
|
slice=expr.index.accept(self),
|
||||||
|
)
|
||||||
|
|
||||||
|
def visit_slice_expr(self, expr: p.SliceExpr) -> ast.expr:
|
||||||
|
return ast.Slice(
|
||||||
|
lower=expr.lower.accept(self) if expr.lower is not None else None,
|
||||||
|
upper=expr.upper.accept(self) if expr.upper is not None else None,
|
||||||
|
step=expr.step.accept(self) if expr.step is not None else None,
|
||||||
|
)
|
||||||
|
|
||||||
|
def visit_expression_stmt(self, stmt: p.ExpressionStmt) -> ast.stmt:
|
||||||
|
return ast.Expr(
|
||||||
|
value=stmt.expr.accept(self),
|
||||||
|
)
|
||||||
|
|
||||||
|
def visit_function(self, stmt: p.Function) -> ast.stmt:
|
||||||
|
return ast.FunctionDef(
|
||||||
|
name=stmt.name,
|
||||||
|
args=ast.arguments(
|
||||||
|
posonlyargs=[ast.arg(arg=arg.name) for arg in stmt.posonlyargs],
|
||||||
|
vararg=None,
|
||||||
|
args=[ast.arg(arg=arg.name) for arg in stmt.args],
|
||||||
|
kwonlyargs=[ast.arg(arg=arg.name) for arg in stmt.kwonlyargs],
|
||||||
|
kwarg=None,
|
||||||
|
defaults=[
|
||||||
|
arg.default.accept(self)
|
||||||
|
for arg in stmt.posonlyargs + stmt.args
|
||||||
|
if arg.default is not None
|
||||||
|
],
|
||||||
|
kw_defaults=[
|
||||||
|
arg.default.accept(self) if arg.default is not None else None
|
||||||
|
for arg in stmt.kwonlyargs
|
||||||
|
],
|
||||||
|
),
|
||||||
|
body=self._visit_body(stmt.body),
|
||||||
|
decorator_list=[],
|
||||||
|
)
|
||||||
|
|
||||||
|
def visit_type_assign(self, stmt: p.TypeAssign) -> ast.stmt:
|
||||||
|
# TODO: is that ok?
|
||||||
|
return ast.Pass()
|
||||||
|
|
||||||
|
def visit_assign_stmt(self, stmt: p.AssignStmt) -> ast.stmt:
|
||||||
|
return ast.Assign(
|
||||||
|
targets=[target.accept(self) for target in stmt.targets],
|
||||||
|
value=stmt.value.accept(self),
|
||||||
|
)
|
||||||
|
|
||||||
|
def visit_return_stmt(self, stmt: p.ReturnStmt) -> ast.stmt:
|
||||||
|
return ast.Return(
|
||||||
|
value=stmt.value.accept(self) if stmt.value is not None else None,
|
||||||
|
)
|
||||||
|
|
||||||
|
def visit_if_stmt(self, stmt: p.IfStmt) -> ast.stmt:
|
||||||
|
return ast.If(
|
||||||
|
test=stmt.test.accept(self),
|
||||||
|
body=self._visit_body(stmt.body),
|
||||||
|
orelse=self._visit_body(stmt.orelse),
|
||||||
|
)
|
||||||
|
|
||||||
|
def _visit_body(self, stmts: list[p.Stmt]) -> list[ast.stmt]:
|
||||||
|
return [stmt.accept(self) for stmt in stmts]
|
||||||
Reference in New Issue
Block a user