MCPcopy
hub / github.com/eth-sri/lmql / transform_node

Method transform_node

src/lmql/language/compiler.py:568–640  ·  view source on GitHub ↗
(self, expr, snf)

Source from the content-addressed store, hash-verified

566 return bn or node.id
567
568 def transform_node(self, expr, snf):
569 if type(expr) is ast.BoolOp:
570 if type(expr.op) is ast.And or type(expr.op) is ast.Or:
571 ops = expr.values
572 tops = [self.transform_node(op, snf) for op in ops]
573 tops_list = ",\n ".join([t.strip() or "None" for t in tops])
574
575 Op = f"{OPS_NAMESPACE}.AndOp" if type(expr.op) is ast.And else f"{OPS_NAMESPACE}.OrOp"
576 return snf.add(f"{Op}([\n {tops_list}\n])")
577 elif type(expr) is ast.Name:
578 return self.transform_name(expr)
579 elif type(expr) is ast.UnaryOp:
580 op = expr.op
581
582 Ops = {
583 ast.Not: f"{OPS_NAMESPACE}.NotOp"
584 }
585
586 for OpT, impl in Ops.items():
587 if type(op) is OpT:
588 operand = self.transform_node(expr.operand, snf).strip()
589 return snf.add(f"{impl}([{operand}])")
590
591 assert False, "unary operator {} not supported.".format(type(expr.op))
592 elif type(expr) is ast.Compare:
593 op = expr.ops[0]
594 assert len(expr.ops) == 1, "compiler currently does not support comparison with more than one operator"
595
596 Ops = {
597 ast.Eq: f"{OPS_NAMESPACE}.EqOp",
598 ast.Lt: f"{OPS_NAMESPACE}.Lt",
599 ast.Gt: f"{OPS_NAMESPACE}.Gt",
600 ast.In: f"{OPS_NAMESPACE}.InOp"
601 }
602
603 for OpT, impl in Ops.items():
604 if type(op) is OpT:
605 ops = [self.transform_node(c, snf) for c in [expr.left] + expr.comparators]
606 ops_list = ", ".join(ops).strip()
607 return snf.add(f"{impl}([{ops_list}])")
608
609 if is_type_constraint(expr):
610 type_name = expr.comparators[0].id
611 var_name = expr.left.args[0].id
612 return snf.add(f"{OPS_NAMESPACE}.CallOp([{LIB_NAMESPACE}.types.is_type, [{OPS_NAMESPACE}.Var('{var_name}'), {type_name}]], locals(), globals())")
613 # return snf.add(f"{LIB_NAMESPACE}.is_type([{type_name}, {OPS_NAMESPACE}.Var('{var_name}')])")
614
615 assert False, "operator {} is not supported.".format(ast.unparse(expr))
616 elif type(expr) is ast.Constant:
617 return self.default_transform_node(expr, snf).strip()
618 elif type(expr) is ast.ListComp:
619 return self.default_transform_node(expr, snf).strip()
620 elif type(expr) is ast.Call:
621 constraint_ref = get_builtin_name(expr.func) or get_inner_constraint_ref(expr.func, self.scope)
622 if constraint_ref is not None:
623 args = [self.transform_node(a, snf) for a in expr.args]
624 args_list = ", ".join(args)
625 return f"{constraint_ref}([{args_list}])"

Callers 3

transformMethod · 0.80
transformMethod · 0.80
visit_AssignMethod · 0.80

Calls 9

transform_nameMethod · 0.95
is_type_constraintFunction · 0.85
get_builtin_nameFunction · 0.85
get_inner_constraint_refFunction · 0.85
stripMethod · 0.80
itemsMethod · 0.80
addMethod · 0.45

Tested by 1

visit_AssignMethod · 0.64