def clean_top_level(tree: ast.Module) -> ast.AST: tree.body = [ node for node in tree.body if isinstance(node, _ast.FunctionDef) \ or isinstance(node, _ast.ClassDef) \ or isinstance(node, _ast.Import) \ or isinstance(node, _ast.ImportFrom) ]
def annotate_funcs_in_module(module: ast.Module, decorator: ast.expr, copy: bool = False) -> ast.Module: if copy: module = deepcopy(module) module.body = [ annotate_funcs_in_stmt(stmt, decorator, copy=False) # no need to copy, since we are already copied the whole module for stmt in module.body ] return module
def visit(self, node: ast.Module) -> ast.Module: """Transform a Python module AST node. :param node: ast.Module object :return: ast.Module object """ node.body = [*self._transform_module_body(node.body)] ast.fix_missing_locations(node) return node
def visit_Module(self, module: ast.Module): """ Visitor of the module node Fills module symbol table :param module: """ mod: Module = Module() self._current_module = mod self._scope_stack.append(SymbolScope()) global_stmts = [] function_stmts = [] for stmt in module.body: if isinstance(stmt, (ast.FunctionDef, ast.AsyncFunctionDef)): function_stmts.append(stmt) elif not (isinstance(stmt, ast.Expr) and isinstance(stmt.value, ast.Constant)): # don't evaluate constant expression - for example: string for documentation if self.visit(stmt) is not Builtin.Event: global_stmts.append(stmt) module.body = global_stmts + function_stmts for var_id, var in mod.variables.items(): # all static fields must be initialized if not mod.is_variable_assigned(var_id): self._log_error( CompilerError.UnresolvedReference( line=var.origin.lineno if var.origin is not None else 0, col=var.origin.col_offset if var.origin is not None else 0, symbol_id=var_id)) for stmt in function_stmts: result = self.visit(stmt) # don't evaluate the metadata function in the following analysers if result is Builtin.Metadata: module.body.remove(stmt) # TODO: include the body of the builtin methods to the ast # TODO: get module name self.modules['main'] = mod module_scope = self._scope_stack.pop() for symbol_id, symbol in module_scope.symbols.items(): if symbol_id in self._global_assigned_variables: mod.include_symbol(symbol_id, symbol) mod.assign_variable(symbol_id) self._global_assigned_variables.clear() self._current_module = None
def create_pdg(tree: ast.Module) -> PDG: tree.body = simplify.simplify_stmts(tree.body) cur_cdg = cdg.create_cdg(tree) ddg_tree = ddg.create_ddg(tree) lookup: t.MutableMapping[t.Union[ast.stmt, ast.ExceptHandler], t.Tuple[cdg.CDGNode, cdg.CDG]] = {} todo = [cur_cdg] while todo: cur = todo.pop() for node in cur.nodes: lookup[node.orig] = (node, cur) todo.extend(cur.subgraphs) for key, vals in ddg_tree.mapping.items(): from_node, _ = lookup[key] for val in vals: to_node, graph = lookup[val] edge = cdg.CDGEdge(from_node, to_node, None, True) graph.add_edge(edge) return PDG(cur_cdg)
def wrap_schema(instrumented_tree: ast.Module): """[summary] This function takes the branch-distance instrumented schema and wraps it around another function such that we can return the schema object instance as python object to run Args: instrumented_tree (ast.Module): AST of schema file with modifications for branch distance already done Returns: ast.Module: AST that upon execution would return the schema as a python executable object """ wrapper = ast.FunctionDef(name='wrapper_function', args=ast.arguments(posonlyargs=[], args=[], vararg=None, kwonlyargs=[], kw_defaults=[], kwarg=None, defaults=[]), decorator_list=[], returns=None, type_comment=None) # We know that the name is "schema" wrapper.body = [*instrumented_tree.body, ast.parse("return schema").body[0]] instrumented_tree.body = [wrapper] ast.fix_missing_locations(instrumented_tree) return instrumented_tree
def apply_fix(self, ctx: common.BuildContext, tree: ast.Module) -> ast.Module: tree.body = self.walk_stmtlist(tree.body) return tree
def __call__(self, cfg: common.BuildConfig, tree: ast.Module) -> ast.Module: tree.body = self.walk_stmtlist(tree.body, tree, "body") return tree