#!/usr/bin/env python3 # -*- coding: utf-8 -*- # # Copyright 2020 Alibaba Group Holding Limited. All Rights Reserved. # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. # You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # """An almost-complete Python to Cython AST transformer, with injected GRAPHSCOPE-specific translation. Python AST nodes are translated to corresponding Cython AST nodes as it is, except: 1. for top-level method, a Cython type annotation is attached to the function signature, for example, .. code:: python @graphscope.analytical.udf.peval('sssp') def PEval(frag, context): ... will be translated as: .. code:: cython cdef public void IncEval(Fragment *frag, ComputeContext *context): ... it will make Cython understand what we really want and generate proper Cpp code. 2. for invokation on methods inside :code:`graphscope.analytical.udf.core`, we generate proper special :code:`cdef` definitions, or proper Cpp invokations, just like :code:`cython.declare`, for example, .. code:: python heap = graphscope.analytical.udf.heap((float, 'node')) modified = lang.vector(bool, [False for _ in range(inner_vertices.size())]) will be translated as: .. code:: cython cdef priority_queue[pair[double, NodeT]] heap cdef vector[bool] modified([False for _ in range(inner_vertices.size())]) Note that :code:`float` in Python is mapped to :code:`double` in Cython (further in Cpp code). More specifically, we define a series of placeholders in module :code:`graphscope.analytical.udf.core`, which cannot be executed in pure python mode. The :code:`graphscope.analytical.udf.xxx` decorators will translate those ordinary *"assignment and call"* into a :code:`cdef` node in Cython AST. """ import ast import copy import functools import inspect import textwrap import types import warnings from Cython.CodeWriter import CodeWriter from Cython.Compiler import Builtin from Cython.Compiler import StringEncoding from Cython.Compiler.ExprNodes import * from Cython.Compiler.ModuleNode import * from Cython.Compiler.Nodes import * from graphscope.analytical.udf.patch import patch_cython_codewriter from graphscope.analytical.udf.utils import CType from graphscope.analytical.udf.utils import ExpectFuncDef from graphscope.analytical.udf.utils import LinesWrapper from graphscope.analytical.udf.utils import PregelAggregatorType from graphscope.analytical.udf.utils import ProgramModel from graphscope.framework.errors import check_argument class GRAPECompiler(ast.NodeVisitor): def __init__(self, name, vd_type, md_type, program_model=ProgramModel.Pregel): """ Args: name: str. The name of class. vd_type: str. The type of the data stored in vertex. md_type: str. The type of the message. program_model: ProgramModel. 'Pregel' or 'PIE' """ self._name = name self._vd_type = vd_type self._md_type = md_type self._program_model = program_model # store aggregate function indexed by name self.__registered_aggregators = {} self.__globals = {} self.__func_params_name_list = [] self.__pyx_header = LinesWrapper() def set_pregel_program_model(self): self._program_model = ProgramModel.Pregel def set_pie_program_model(self): self._program_model = ProgramModel.PIE def parse(self, source): """Parse source into cython module node object. source: str The source code may represent a statement or expression. Raises: RuntimeError: unsupported ast trans from python to cython. """ tree = ast.parse(textwrap.dedent(source)) # associate `parent` reference to every node for node in ast.walk(tree): for child in ast.iter_child_nodes(node): setattr(child, "__parent__", node) cyast = self.visit(tree) return cyast def run(self, func_or_ast, pyx_header): self.__pyx_header = pyx_header # we already has a AST: just run it if isinstance(func_or_ast, ast.AST): cyast = self.visit(func_or_ast) else: check_argument(isinstance(func_or_ast, types.FunctionType)) # ignore varargs and keywords self.__func_params_name_list = inspect.getfullargspec(func_or_ast).args self.__globals = func_or_ast.__globals__ cyast = self.parse(inspect.getsource(func_or_ast)) writer = patch_cython_codewriter(CodeWriter()) cycode = "\n".join(writer.write(cyast).lines) return cycode def compile(self, source): """Compile source into cython code.""" cyast = self.parse(source) writer = patch_cython_codewriter(CodeWriter()) return "\n".join(writer.write(cyast).lines) def make_plain_arg(self, name, arg_loc): return CArgDeclNode( arg_loc, base_type=CSimpleBaseTypeNode( arg_loc, name=None, is_basic_c_type=0, signed=1, longness=0, is_self_arg=False, ), declarator=CNameDeclaratorNode(arg_loc, name=name), not_none=0, or_none=0, default=None, annotation=None, ) def make_value_arg(self, value_type, name, arg_loc): return CArgDeclNode( arg_loc, base_type=CSimpleBaseTypeNode( arg_loc, name=value_type, is_basic_c_type=0, signed=1, longness=0, is_self_arg=False, ), declarator=CNameDeclaratorNode(arg_loc, name=name), not_none=0, or_none=0, default=None, annotation=None, ) def make_ptr_arg(self, ptr_type, name, arg_loc): return CArgDeclNode( arg_loc, base_type=CSimpleBaseTypeNode( arg_loc, name=ptr_type, is_basic_c_type=0, signed=1, longness=0, is_self_arg=False, ), declarator=CPtrDeclaratorNode( arg_loc, base=CNameDeclaratorNode(arg_loc, name=name) ), not_none=0, or_none=0, default=None, annotation=None, ) def make_ref_arg(self, ref_type, name, arg_loc): return CArgDeclNode( arg_loc, base_type=CSimpleBaseTypeNode( arg_loc, name=ref_type, is_basic_c_type=0, signed=1, longness=0, complex=0, is_self_arg=False, templates=None, ), declarator=CReferenceDeclaratorNode( arg_loc, base=CNameDeclaratorNode(arg_loc, name=name) ), not_none=0, or_none=0, default=None, annotation=None, ) def make_template_arg( self, value_type, value_tpls, name, arg_loc, use_ptr=False, use_ref=False ): def mk_tpl_arg(n): return CComplexBaseTypeNode( arg_loc, base_type=CSimpleBaseTypeNode( arg_loc, name=n, is_basic_c_type=0, signed=1, longness=0, is_self_arg=False, ), declarator=CNameDeclaratorNode( arg_loc, name="", cname=None, default=None ), ) tpl_type = TemplatedTypeNode( arg_loc, positional_args=[mk_tpl_arg(n) for n in value_tpls], keyword_args=DictNode(arg_loc, key_value_pairs=[]), base_type_node=CSimpleBaseTypeNode( arg_loc, name=value_type, is_basic_c_type=0, signed=1, longness=0, is_self_arg=False, ), ) if use_ptr: declarator = CPtrDeclaratorNode( arg_loc, base=CNameDeclaratorNode(arg_loc, name=name) ) elif use_ref: declarator = CReferenceDeclaratorNode( arg_loc, base=CNameDeclaratorNode(arg_loc, name=name) ) else: declarator = CNameDeclaratorNode(arg_loc, name=name) return CArgDeclNode( arg_loc, base_type=tpl_type, declarator=declarator, not_none=0, or_none=0, default=None, annotation=None, ) def loc(self, node): return ["", 0, 0] def generic_visit(self, node): raise NotImplementedError("AST node %s is not supported yet" % node) def visit_Module(self, node): body = self.visit(node.body[0]) return ModuleNode(self.loc(node), body=body) def visit_ImportFrom(self, node): raise RuntimeError("ImportFrom is not supported yet.") def visit_Import(self, node): raise RuntimeError("Import is not supported yet.") def visit_ClassDef(self, node): raise RuntimeError("Class definition is not supported yet.") def visit_JoinedStr(self, node): raise RuntimeError("Joinedstr is not supported yet.") def visit_Constant(self, node): if isinstance(node.value, int): return IntNode(self.loc(node), value=str(node.value)) if isinstance(node.value, float): # We won't have c float, we map floating types to double return FloatNode(self.loc(node), value=str(node.value)) if isinstance(node.value, str): if node.kind == "u": return UnicodeNode(self.loc(node), value=node.value, bytes_value=None) return StringNode( self.loc(node), value=node.value, unicode_value=StringEncoding.EncodedString(node.value), ) if ( isinstance(Ellipsis, type) and isinstance(node.value, Ellipsis) or isinstance(node.value, type(Ellipsis)) ): return EllipsisNode(self.loc(node)) if isinstance(node.value, bytes): return BytesNode(self.loc(node), value=node.s) if node.value is None: return NoneNode(self.loc(node)) raise NotImplementedError("Unknown constant value: %s" % node) def visit_Num(self, node): if isinstance(node.n, int): return IntNode(self.loc(node), value=str(node.n)) if isinstance(node.n, float): return FloatNode(self.loc(node), value=str(node.n)) if isinstance(node.n, complex): raise NotImplementedError("Not support complex constant yet") raise NotImplementedError("Unknown constant value: %s" % node) def visit_Str(self, node): return StringNode( self.loc(node), value=node.s, unicode_value=StringEncoding.EncodedString(node.s), ) def visit_Bytes(self, node): return BytesNode(self.loc(node), value=node.s) def visit_List(self, node): return ListNode(self.loc(node), args=[self.visit(elt) for elt in node.elts]) def visit_Tuple(self, node): return TupleNode(self.loc(node), args=[self.visit(elt) for elt in node.elts]) def visit_Set(self, node): return SetNode(self.loc(node), args=[self.visit(elt) for elt in node.elts]) def visit_Dict(self, node): kvs = [ DictItemNode(self.loc(node), key=self.visit(k), value=self.visit(v)) for k, v in zip(node.keys, node.values) ] return DictNode(self.loc(node), key_value_pairs=kvs) def visit_Ellipsis(self, node): return EllipsisNode(self.loc(node)) def visit_NameConstant(self, node): if node.value in [True, False]: return BoolNode(self.loc(node), value=node.value) return NoneNode(self.loc(node)) def visit_Name(self, node): return NameNode(self.loc(node), name=node.id) def visit_Expr(self, node): expr = self.visit(node.value) if isinstance(expr, CVarDefNode): return expr return ExprStatNode(self.loc(node), expr=expr) def visit_UnaryOp(self, node): if isinstance(node.op, ast.UAdd): return UnaryPlusNode( self.loc(node), operator="+", operand=self.visit(node.operand) ) if isinstance(node.op, ast.USub): return UnaryMinusNode( self.loc(node), operator="-", operand=self.visit(node.operand) ) if isinstance(node.op, ast.Not): return NotNode(self.loc(node), operand=self.visit(node.operand)) if isinstance(node.op, ast.Invert): return TildeNode( self.loc(node), operator="~", operand=self.visit(node.operand) ) def visit_UAdd(self, node): return "+" def visit_USub(self, node): return "-" def visit_Not(self, node): return "not" def visit_Invert(self, node): return "invert" def visit_BinOp(self, node): lhs = self.visit(node.left) rhs = self.visit(node.right) op_mapping = { ast.Add: (AddNode, "+"), ast.Sub: (SubNode, "-"), ast.Mult: (MulNode, "*"), ast.Div: (DivNode, "/"), ast.FloorDiv: (DivNode, "//"), ast.Mod: (ModNode, "%"), ast.Pow: (PowNode, "**"), ast.MatMult: (MatMultNode, "@"), ast.LShift: (IntBinopNode, "<<"), ast.RShift: (IntBinopNode, ">>"), ast.BitOr: (IntBinopNode, "|"), ast.BitXor: (IntBinopNode, "^"), ast.BitAnd: (IntBinopNode, "&"), } op_type, op = op_mapping[type(node.op)] return op_type(self.loc(op), operator=op, operand1=lhs, operand2=rhs) def visit_Add(self, node): return "+" def visit_Sub(self, node): return "-" def visit_Mult(self, node): return "*" def visit_Div(self, node): return "/" def visit_FloorDiv(self, node): return "//" def visit_Mod(self, node): return "%" def visit_Pow(self, node): return "**" def visit_LShift(self, node): return "<<" def visit_RShift(self, node): return ">>" def visit_BitOr(self, node): return "|" def visit_BitXor(self, node): return "^" def visit_BitAnd(self, node): return "&" def visit_MatMult(self, node): return "@" def visit_AnnAssign(self, node): annotation = NameNode(self.loc(node), name=node.annotation.id) lhs = NameNode(self.loc(node), name=node.target.id, annotation=annotation) rhs = self.visit(node.value) return SingleAssignmentNode(self.loc(node), lhs=lhs, rhs=rhs) def visit_BoolOp(self, node): return BoolBinopNode( self.loc(node), operator=self.visit(node.op), operand1=self.visit(node.values[0]), operand2=self.visit(node.values[1]), ) def visit_And(self, node): return "and" def visit_Or(self, node): return "or" def visit_Compare(self, node): operator = self.visit(node.ops[0]) operand1 = self.visit(node.left) operand2 = self.visit(node.comparators[0]) if len(node.comparators) == 1: # single comparison return PrimaryCmpNode( self.loc(node), operator=operator, operand1=operand1, operand2=operand2 ) # multiple continuous comparison cascade_node = CascadedCmpNode( self.loc(node), operator=self.visit(node.ops[-1]), operand2=self.visit(node.comparators[-1]), ) for op, comparator in zip(node.ops[-2:0:-1], node.comparators[-2:0:-1]): cascade_node = CascadedCmpNode( self.loc(node), operator=self.visit(op), operand2=self.visit(comparator), cascade=cascade_node, ) return PrimaryCmpNode( self.loc(node), operator=operator, operand1=operand1, operand2=operand2, cascade=cascade_node, ) def visit_Eq(self, node): return "==" def visit_NotEq(self, node): return "!=" def visit_Lt(self, node): return "<" def visit_LtE(self, node): return "<=" def visit_Gt(self, node): return ">" def visit_GtE(self, node): return ">=" def visit_Is(self, node): return "is" def visit_IsNot(self, node): return "is not" def visit_In(self, node): return "in" def visit_NotIn(self, node): return "not in" def __flatten_func_name(self, name): if isinstance(name, ast.Name): return [name.id] if isinstance(name, ast.Attribute): return self.__flatten_func_name(name.value) + [name.attr] return [] def __is_graphscope_api_call(self, node): flat_func_name = self.__flatten_func_name(node.func) if len(flat_func_name) == 0: return False if flat_func_name[0] in self.__func_params_name_list: return True # check from graphscope module cascade = self.__globals.get(flat_func_name[0]) if cascade is None: return False for n in flat_func_name[1:]: if cascade is None or not hasattr(cascade, n): return False cascade = getattr(cascade, n) return cascade.__module__ == "graphscope.analytical.udf.types" def __visit_GraphScopeAPICall(self, node): full_func_name = self.__flatten_func_name(node.func) obj = full_func_name[0] name = node.func.attr if obj == "graphscope": # graphscope.declare() if name == "declare": var = node.args[1].id var_type = node.args[0].attr return CVarDefNode( self.loc(node), base_type=CSimpleBaseTypeNode( self.loc(node), name=var_type, module_path=[], is_basic_c_type=0, signed=1, ), declarators=[CNameDeclaratorNode(self.loc(node), name=var)], visibility="private", ) elif obj == "context" and name == "register_aggregator": # context.register_aggregator() args = node.args if len(args) != 2: raise ValueError("Params within register_aggregator incorrect.") if ( isinstance(args[1], ast.Attribute) and args[1].value.id == "PregelAggregatorType" ): self.__registered_aggregators[str(args[0].s)] = args[1].attr return SimpleCallNode( self.loc(node), function=self.visit(node.func), args=[self.visit(arg) for arg in node.args], ) elif obj == "context" and name == "aggregate": # context.aggregate() args = node.args if len(args) != 2: raise ValueError("Params within aggregate incorrect.") if str(args[0].s) not in self.__registered_aggregators.keys(): raise KeyError( "Aggregator %s not exist, you may want to register first." % str(args[0].s) ) ctype = PregelAggregatorType.extract_ctype( self.__registered_aggregators[str(args[0].s)] ) return SimpleCallNode( self.loc(node), function=IndexNode( self.loc(node), base=AttributeNode( self.loc(node), obj=NameNode(self.loc(node), name=obj), attribute=name, ), index=NameNode(self.loc(node), name=str(ctype)), ), args=[self.visit(arg) for arg in node.args], ) elif obj == "context" and name == "get_aggregated_value": # context.get_aggregated_value() args = node.args if len(args) != 1: raise ValueError("Params within get_aggregated_value incorrect.") if str(args[0].s) not in self.__registered_aggregators.keys(): raise KeyError( "Aggregator %s not exist, you may want to register first." % str(args[0].s) ) ctype = PregelAggregatorType.extract_ctype( self.__registered_aggregators[str(args[0].s)] ) return SimpleCallNode( self.loc(node), function=IndexNode( self.loc(node), base=AttributeNode( self.loc(node), obj=NameNode(self.loc(node), name=obj), attribute=name, ), index=NameNode(self.loc(node), name=str(ctype)), ), args=[self.visit(arg) for arg in node.args], ) elif obj == "context" and full_func_name[1] == "math": mnode = copy.copy(node) mnode.func = ast.Attribute(value=ast.Name(id="math"), attr=name) return self.visit(mnode) else: return SimpleCallNode( self.loc(node), function=self.visit(node.func), args=[self.visit(arg) for arg in node.args], ) def visit_Call(self, node): if self.__is_graphscope_api_call(node): return self.__visit_GraphScopeAPICall(node) if not node.keywords: return SimpleCallNode( self.loc(node), function=self.visit(node.func), args=[self.visit(arg) for arg in node.args], ) # with kwargs param return GeneralCallNode( self.loc(node), function=self.visit(node.func), positional_args=TupleNode( self.loc(node), args=[self.visit(arg) for arg in node.args] ), # keyword_args=DictNode(self.loc(node), key_value_pairs=[])) keyword_args=self._visit_keywords(node.keywords), ) def _visit_keywords(self, node): kvs = [] for keyword in node: kvs.append(self.visit_keyword(keyword)) return DictNode(self.loc(node), key_value_pairs=kvs, reject_duplicates=True) def visit_keyword(self, node): key = IdentifierStringNode(self.loc(node), value=node.arg) return DictItemNode(self.loc(node), key=key, value=self.visit(node.value)) def visit_IfExp(self, node): return CondExprNode( self.loc(node), test=self.visit(node.test), true_val=self.visit(node.body), false_val=self.visit(node.orelse), ) def visit_Attribute(self, node): full_attr_name = self.__flatten_func_name(node) if full_attr_name[0] == "context" and full_attr_name[1] == "math": mnode = copy.copy(node) mnode.value = ast.Name(id="math") return self.visit(mnode) return AttributeNode( self.loc(node), obj=self.visit(node.value), attribute=node.attr ) def visit_Subscript(self, node): return IndexNode( self.loc(node), base=self.visit(node.value), index=self.visit(node.slice) ) def visit_Index(self, node): return self.visit(node.value) def visit_Slice(self, node): start = ( NoneNode(self.loc(node)) if node.lower is None else self.visit(node.lower) ) stop = ( NoneNode(self.loc(node)) if node.upper is None else self.visit(node.upper) ) step = NoneNode(self.loc(node)) if node.step is None else self.visit(node.step) return SliceNode(self.loc(node), start=start, stop=stop, step=step) def visit_ExtSlice(self, node): return TupleNode(self.loc(node), args=[self.visit(dim) for dim in node.dims]) def visit_ListComp(self, node): check_argument(len(node.generators) == 1) # has if node or not has_if = True if node.generators[0].ifs else False expression_value = self.visit(node.elt) generator = node.generators[0] iter_value = IteratorNode( self.loc(generator.iter), sequence=self.visit(generator.iter) ) comp_node = ComprehensionAppendNode(self.loc(generator), expr=expression_value) if has_if: check_argument(len(node.generators[0].ifs) == 1) # construct IfStatNode condition = self.visit(node.generators[0].ifs[0]) body = comp_node if_stat_node = IfStatNode( self.loc(node), if_clauses=[ IfClauseNode(self.loc(node), condition=condition, body=body) ], else_clause=None, ) loop = ForInStatNode( self.loc(node), target=self.visit(generator.target), iterator=iter_value, body=if_stat_node, else_clause=None, is_async=False, ) else: loop = ForInStatNode( self.loc(node), target=self.visit(generator.target), iterator=iter_value, body=comp_node, else_clause=None, is_async=False, ) return ComprehensionNode( self.loc(node), loop=loop, append=comp_node, type=Builtin.list_type, has_local_scope=True, ) def visit_SetComp(self, node): assert len(node.generators) == 1 # has if node or not has_if = True if node.generators[0].ifs else False expression_value = self.visit(node.elt) generator = node.generators[0] iter_value = IteratorNode( self.loc(generator.iter), sequence=self.visit(generator.iter) ) comp_node = ComprehensionAppendNode(self.loc(generator), expr=expression_value) if has_if: assert len(node.generators[0].ifs) == 1 # construct IfStatNode condition = self.visit(node.generators[0].ifs[0]) body = comp_node if_stat_node = IfStatNode( self.loc(node), if_clauses=[ IfClauseNode(self.loc(node), condition=condition, body=body) ], else_clause=None, ) loop = ForInStatNode( self.loc(node), target=self.visit(generator.target), iterator=iter_value, body=if_stat_node, else_clause=None, is_async=False, ) else: loop = ForInStatNode( self.loc(node), target=self.visit(generator.target), iterator=iter_value, body=comp_node, else_clause=None, is_async=False, ) return ComprehensionNode( self.loc(node), loop=loop, append=comp_node, type=Builtin.set_type ) def visit_DictComp(self, node): assert len(node.generators) == 1 # has if node or not has_if = True if node.generators[0].ifs else False generator = node.generators[0] iter_value = IteratorNode( self.loc(generator.iter), sequence=self.visit(generator.iter) ) comp_node = DictComprehensionAppendNode( self.loc(generator), key_expr=self.visit(node.key), value_expr=self.visit(node.value), ) if has_if: assert len(node.generators[0].ifs) == 1 # construct IfStatNode condition = self.visit(node.generators[0].ifs[0]) body = comp_node if_stat_node = IfStatNode( self.loc(node), if_clauses=[ IfClauseNode(self.loc(node), condition=condition, body=body) ], else_clause=None, ) loop = ForInStatNode( self.loc(node), target=self.visit(generator.target), iterator=iter_value, body=if_stat_node, else_clause=None, is_async=False, ) else: loop = ForInStatNode( self.loc(node), target=self.visit(generator.target), iterator=iter_value, body=comp_node, else_clause=None, is_async=False, ) return ComprehensionNode( self.loc(node), loop=loop, append=comp_node, type=Builtin.dict_type ) def visit_Assign(self, node): # `tuple` represents a multiple assign assert len(node.targets) == 1 if ( hasattr(node.targets[0], "id") and node.targets[0].id in self.__func_params_name_list ): raise RuntimeError("Can't assign to internal variables.") lhs = self.visit(node.targets[0]) rhs = self.visit(node.value) return SingleAssignmentNode(self.loc(node), lhs=lhs, rhs=rhs) def visit_AugAssign(self, node): return InPlaceAssignmentNode( self.loc(node), operator=self.visit(node.op), lhs=self.visit(node.target), rhs=self.visit(node.value), ) def visit_Raise(self, node): return RaiseStatNode( self.loc(node), exc_type=self.visit(node.exc), exc_value=None, exc_tb=None, cause=None if node.cause is None else self.visit(node.cause), ) def visit_ExceptHandler(self, node): if node.type: pattern = [self.visit(node.type)] if node.name: target = NameNode(self.loc(node), name=node.name) else: target = None else: pattern = None target = None body = StatListNode( self.loc(node), stats=[self.visit(stat) for stat in node.body] ) return ExceptClauseNode( self.loc(node), pattern=pattern, target=target, body=body, is_except_as=False, ) def visit_Try(self, node): body = StatListNode( self.loc(node), stats=[self.visit(stat) for stat in node.body] ) except_clauses = [self.visit(ec) for ec in node.handlers] if node.orelse: else_clause = StatListNode( self.loc(node), stats=[self.visit(stat) for stat in node.orelse] ) else: else_clause = None try_except_stat_node = TryExceptStatNode( self.loc(node), body=body, except_clauses=except_clauses, else_clause=else_clause, ) # with `finally` statement or not if node.finalbody: final_clause = StatListNode( self.loc(node), stats=[self.visit(stat) for stat in node.finalbody] ) return TryFinallyStatNode( self.loc(node), body=try_except_stat_node, finally_clause=final_clause ) return try_except_stat_node def visit_Assert(self, node): return AssertStatNode( self.loc(node), cond=self.visit(node.test), value=self.visit(node.msg) if node.msg else None, ) def visit_Delete(self, node): return DelStatNode( self.loc(node), args=[self.visit(target) for target in node.targets] ) def visit_Pass(self, node): return PassStatNode(self.loc(node)) def visit_If(self, node): condition = self.visit(node.test) body = StatListNode( self.loc(node), stats=[self.visit(stat) for stat in node.body] ) if node.orelse: else_body = StatListNode( self.loc(node), stats=[self.visit(stat) for stat in node.orelse] ) else: else_body = None return IfStatNode( self.loc(node), if_clauses=[IfClauseNode(self.loc(node), condition=condition, body=body)], else_clause=else_body, ) def visit_For(self, node): target_value = self.visit(node.target) iter_value = IteratorNode(self.loc(node.iter), sequence=self.visit(node.iter)) body = StatListNode( self.loc(node), stats=[self.visit(stat) for stat in node.body] ) if node.orelse: else_body = StatListNode( self.loc(node), stats=[self.visit(stat) for stat in node.orelse] ) else: else_body = None return ForInStatNode( self.loc(node), target=target_value, iterator=iter_value, body=body, else_clause=else_body, is_async=False, ) def visit_While(self, node): condition = self.visit(node.test) body = StatListNode( self.loc(node), stats=[self.visit(stat) for stat in node.body] ) if node.orelse: else_body = StatListNode( self.loc(node), stats=[self.visit(stat) for stat in node.orelse] ) else: else_body = None return WhileStatNode( self.loc(node), condition=condition, body=body, else_clause=else_body ) def visit_withitem(self, node): return self.visit(node.context_expr) def visit_With(self, node): # multiple items is not supported yet assert len(node.items) == 1 manager = self.visit(node.items[0]) target = self.visit(node.items[0].optional_vars) body = StatListNode( self.loc(node), stats=[self.visit(stat) for stat in node.body] ) return WithStatNode(self.loc(node), manager=manager, target=target, body=body) def visit_Break(self, node): return BreakStatNode(self.loc(node)) def visit_Continue(self, node): return ContinueStatNode(self.loc(node)) def visit_FunctionDef(self, node): def is_static_method(func): return ( func.decorator_list and isinstance(func.decorator_list[0], ast.Name) and (func.decorator_list[0].id == "staticmethod") ) if not is_static_method(node): raise RuntimeError("Missing decorator staticmethod.") function_name = node.name function_return_type = "void" if self._program_model == ProgramModel.PIE: # PIE program model if function_name == ExpectFuncDef.INIT.value: args = node.args.args assert len(args) == 2, "The number of parameters does not match" args = [ self.make_ref_arg("Fragment", args[0].arg, self.loc(args[0])), self.make_template_arg( "Context", [self._vd_type, self._md_type], args[1].arg, self.loc(args[1]), use_ref=True, ), ] elif function_name == ExpectFuncDef.PEVAL.value: args = node.args.args assert len(args) == 2, "The number of parameters does not match" args = [ self.make_ref_arg("Fragment", args[0].arg, self.loc(args[0])), self.make_template_arg( "Context", [self._vd_type, self._md_type], args[1].arg, self.loc(args[1]), use_ref=True, ), ] elif function_name == ExpectFuncDef.INCEVAL.value: args = node.args.args assert len(args) == 2, "The number of parameters does not match" args = [ self.make_ref_arg("Fragment", args[0].arg, self.loc(args[0])), self.make_template_arg( "Context", [self._vd_type, self._md_type], args[1].arg, self.loc(args[1]), use_ref=True, ), ] else: raise RuntimeError( "Not recognized method named {}".format(function_name) ) elif self._program_model == ProgramModel.Pregel: if function_name == ExpectFuncDef.INIT.value: args = node.args.args assert len(args) == 2, "The number of parameters does not match" args = [ self.make_template_arg( "Vertex", [self._vd_type, self._md_type], args[0].arg, self.loc(args[0]), use_ref=True, ), self.make_template_arg( "Context", [self._vd_type, self._md_type], args[1].arg, self.loc(args[1]), use_ref=True, ), ] elif function_name == ExpectFuncDef.COMPUTE.value: args = node.args.args assert len(args) == 3, "The number of parameters does not match" args = [ self.make_template_arg( "MessageIterator", [self._md_type], args[0].arg, self.loc(args[0]), ), self.make_template_arg( "Vertex", [self._vd_type, self._md_type], args[1].arg, self.loc(args[1]), use_ref=True, ), self.make_template_arg( "Context", [self._vd_type, self._md_type], args[2].arg, self.loc(args[2]), use_ref=True, ), ] elif function_name == ExpectFuncDef.COMBINE.value: args = node.args.args assert len(args) == 1, "The number of parameters does not match" args = [ self.make_template_arg( "MessageIterator", [self._md_type], args[0].arg, self.loc(args[0]), ) ] function_return_type = self._md_type else: raise RuntimeError( "Not recognized method named {}".format(function_name) ) base_type = CSimpleBaseTypeNode( self.loc(node), name=function_return_type, is_basic_c_type=1, signed=1, longness=0, is_self_arg=False, ) declarator_name = function_name declarator = CFuncDeclaratorNode( self.loc(node), base=CNameDeclaratorNode(self.loc(node), name=declarator_name), args=args, has_varargs=False, exception_value=None, exception_check=False, nogil=True, with_gil=False, overridable=False, ) # traverse body body = StatListNode( self.loc(node), stats=[self.visit(expr) for expr in node.body] ) return CFuncDefNode( self.loc(node), visibility="public", base_type=base_type, declarator=declarator, body=body, modifiers=[], api=False, overridable=False, is_const_method=False, ) def visit_Lambda(self, node): return LambdaNode( self.loc(node), args=[ self.make_plain_arg(arg.arg, self.loc(arg)) for arg in node.args.args ], star_arg=None, starstar_arg=None, retult_expr=self.visit(node.body), ) def visit_Return(self, node): if node.value is None: value = None else: value = self.visit(node.value) return ReturnStatNode(self.loc(node), value=value) def visit_Yield(self, node): return YieldExprNode(self.loc(node), expr=self.visit(node.value)) def visit_YieldFrom(self, node): return YieldFromExprNode(self.loc(node), expr=self.visit(node.value)) def visit_Global(self, node): return GlobalNode(self.loc(node), names=node.names) def visit_Nonlocal(self, node): return NonlocalNode(self.loc(node), names=node.names) def visit_Await(self, node): return AwaitExprNode(self.loc(node), expr=self.visit(node.value))