def find_def_code(variables: set, module: dict, ancestors, udc, duc): nodes_def_vars = [ find_last_def_node(variable, module) for variable in variables ] nodes_def_vars.sort(key=lambda x: x.lineno) capturex = CaptureX( list(nodes_def_vars), module, ancestors, defuse_chains=duc, usedef_chains=udc, ) lines_ext = [] for node in capturex.external: lines_ext.append(extast.unparse(node).strip()) for node in nodes_def_vars: line = extast.unparse(node).strip() if line not in lines_ext: lines_ext.append(line) return "\n".join(lines_ext)
def get_annotations(object_def, namespace): """Create the annotations from a definition node""" # print_dump(object_def) ast_annotations = ast.Assign( targets=[extast.Name("annotations", ast.Store())], value=ast.Dict(keys=[], values=[]), type_comment=None, ) if isinstance(object_def, ast.FunctionDef): _fill_ast_annotations_function(object_def, ast_annotations) elif isinstance(object_def, ast.ClassDef): _fill_ast_annotations_class(object_def, ast_annotations) else: raise NotImplementedError # print_dump(ast_annotations) source = extast.unparse(ast_annotations) try: del namespace["__builtins__"] except KeyError: pass exec(source, namespace) return namespace["annotations"]
def _make_code_blocks(self, blocks): code = [] signatures_blocks = [] for block in blocks: str_variables = ", ".join(block.signatures[0].keys()) fdef_block = extast.gast.parse( f"""def {block.name}({str_variables}):pass""").body[0] # TODO: locals_types for blocks locals_types = None signatures_blocks.extend( self._make_header_from_fdef_annotations( fdef_block, block.signatures, locals_types)) code.append(f"\ndef {block.name}({str_variables}):\n") code.append(indent(extast.unparse(block.ast_code), " ")) if block.results: code.append(f" return {', '.join(block.results)}\n") arguments_blocks = { block.name: list(block.signatures[0].keys()) for block in blocks } if arguments_blocks: self._append_line_header_variable(signatures_blocks, "arguments_blocks") code.append(f"arguments_blocks = {str(arguments_blocks)}\n") return signatures_blocks, code
def _make_code_from_fdef_node(self, fdef): if hasattr(fdef, "_transonic_keywords"): decorator_keywords = fdef._transonic_keywords else: decorator_keywords = {} parts = [] if not decorator_keywords.get("boundscheck", True): parts.append("@cython.boundscheck(False)") if not decorator_keywords.get("wraparound", True): parts.append("@cython.wraparound(False)") if decorator_keywords.get("cdivision", False): parts.append("@cython.cdivision(True)") if not decorator_keywords.get("nonecheck", True): parts.append("@cython.noneckeck(False)") if decorator_keywords.get("nogil", False): parts.append("@cython.nogil") transformed = TypeHintRemover().visit(fdef) # convert the AST back to source code parts.append(unparse(transformed)) return format_str("\n".join(parts))
def filter_code_typevars(module, duc, ancestors): """Create a filtered code with what is needed to create the annotations""" module_filtered = ast.Module() kept = module_filtered.body = [] module_filtered.type_ignores = [] suppressed = set() def fill_suppressed(def_): for user in def_.users(): parent_in_body = ancestors.parents(user.node)[1] suppressed.add(parent_in_body) fill_suppressed(user) for node in module.body: if node in suppressed: continue if isinstance(node, ast.Import): if node.names[0].name in ["transonic", "numpy"]: kept.append(node) else: def_ = duc.chains[node.names[0]] fill_suppressed(def_) # suppressed.add() elif isinstance(node, ast.ImportFrom): if node.module in ["transonic", "numpy"]: kept.append(node) elif isinstance(node, (ast.Assign, ast.AugAssign)): kept.append(node) return extast.unparse(module_filtered)
def change_import_name(code_dep: str, changed_node: object, func_name: str, relative: str = None): """Change the name of changed_node in code_dep by adding "__" + func + "__" at the beginning of the imported module, and return the modified code """ mod = extast.parse(code_dep) for node in mod.body: if extast.unparse(node) == extast.unparse(changed_node): if isinstance(node, ast.ImportFrom): node.module = f"__ext__{func_name}__{node.module}" elif isinstance(node, ast.Import): node.names[ 0].name = f"__ext__{func_name}__{node.names[0].name}" if not relative: node.level = 0 return extast.unparse(mod)
def extract_variable_annotations(fdef, namespace): variable_types = {} for node in fdef.body: if isinstance(node, ast.AnnAssign): name = node.target.id source = extast.unparse(node.annotation) result = eval(source, namespace) variable_types[name] = result return variable_types
def strip_typehints(source): """Strip the type hints from a function""" source = format_str(source) # parse the source code into an AST parsed_source = ast.parse(source) # remove all type annotations, function return type definitions # and import statements from 'typing' transformed = TypeHintRemover().visit(parsed_source) # convert the AST back to source code striped_code = extast.unparse(transformed) return striped_code
def add_numba_comments(code): """Add Numba code in Python comments""" mod = parse(code) new_body = [CommentLine("# __protected__ from numba import njit")] for node in mod.body: if isinstance(node, gast.FunctionDef): new_body.append(CommentLine("# __protected__ @njit(cache=True, fastmath=True)")) new_body.append(node) mod.body = new_body return format_str(unparse(mod))
def filter_external_code(module: object, names: list): """Filter the module to keep only the necessary nodes needed by functions or class in the parameter names """ code_dependance_annotations = "" lines_code = [] for node in module.body: for name in names: if isinstance(node, ast.FunctionDef): if node.name == extast.unparse(name).rstrip("\n\r").strip(): ancestors = beniget.Ancestors() ancestors.visit(module) duc = beniget.DefUseChains() duc.visit(module) udc = beniget.UseDefChains(duc) capturex = CaptureX( [node], module, ancestors, defuse_chains=duc, usedef_chains=udc, consider_annotations=None, ) lines_code.append(str(extast.unparse(node))) code_dependance_annotations = capturex.make_code_external() if isinstance(node, ast.Assign): if (node.targets[0].id == extast.unparse(name).rstrip( "\n\r").strip()): lines_code.append(str(extast.unparse(node))) if isinstance(node, ast.ClassDef): if node.name == extast.unparse(name).rstrip("\n\r").strip(): lines_code.append(str(extast.unparse(node))) return code_dependance_annotations + "\n" + "\n".join(lines_code)
def make_backend_source(self, info_analysis, func, path_backend): func_name = func.__name__ jitted_dicts = info_analysis["jitted_dicts"] src = info_analysis["codes_dependance"][func_name] if func_name in info_analysis["special"]: if func_name in jitted_dicts["functions"]: src += "\n" + extast.unparse( jitted_dicts["functions"][func_name]) elif func_name in jitted_dicts["methods"]: src += "\n" + extast.unparse( jitted_dicts["methods"][func_name]) else: # TODO find a prettier solution to remove decorator for cython # than doing two times a regex src += "\n" + re.sub(r"@.*?\sdef\s", "def ", get_source_without_decorator(func)) has_to_write = True if path_backend.exists() and mpi.rank == 0: with open(path_backend) as file: src_old = file.read() if src_old == src: has_to_write = False return src, has_to_write
def make_code_from_fdef_node(fdef): transformed = TypeHintRemover().visit(fdef) # convert the AST back to source code code = extast.unparse(transformed) return format_str(code)
def adapt_code_dependance(func: str, codes_dependance: str, jitted_dicts: dict): """ Adapt code_dependance to the call of a jitted function in a jitted function: - Remove the import transonic - Remove the jitted function statement (i.e func = jit(func)) - Add a import statement to the jitted function - remove the definition of the jitted function if its on the file, or remove the import statement """ special = [] module = extast.parse(codes_dependance) module_body = module.body.copy() jitted_functions = [] for node in module_body: # remove the transonic import if isinstance(node, ast.ImportFrom): if node.module == "transonic": module.body.remove(node) # remove the jitted function by jit() (i.e. func = jit(func)) # and add the import statement for the jitted function elif (isinstance(node, ast.Assign) and isinstance(node.value, ast.Call) and node.value.func.id == "jit"): # if assigned name different from jitted function name if node.targets[0].id != node.value.args[0].id: # change function names in jitted dict # The list "special" contains functions that has to be write from jitted_dicts # see transonic.justintime:287 def_func = extast.unparse(jitted_dicts["functions"][func]) spl = re.split(r"(\W+)", def_func) spl = [ node.value.args[0].id if x == node.targets[0].id else x for x in spl ] st = "".join(str(e) for e in spl) jitted_dicts["functions"][func] = extast.parse(st) special.append(func) jitted_functions.append(node.value.args[0].id) else: jitted_functions.append(node.targets[0].id) module.body.remove(node) module.body.insert( 0, [ extast.parse("from " + node.value.args[0].id + " import " + node.value.args[0].id) ], ) # remove the definition: for node in module_body: # of the jitted function in the file if isinstance(node, ast.FunctionDef): if node.name in jitted_functions: module.body.remove(node) # of the jitted imported function if isinstance(node, ast.ImportFrom): for name in node.names: if name.name in jitted_functions: node.names.remove(name) # remove the importFrom if no function is imported if not node.names: module.body.remove(node) return extast.unparse(module), jitted_dicts, special, jitted_functions
def make_code_external(self): code = [] for node in self.external: code.append(extast.unparse(node).strip()) return "\n".join(code)
def print_unparsed(node): """Print the code corresponding to a tree or a node""" print(extast.unparse(node))
def _make_code_method(self, class_name, fdef, meth_name, annotations, boosted_dicts): class_def = boosted_dicts["classes"][class_name] if class_name in annotations["classes"]: annotations_class = annotations["classes"][class_name] else: annotations_class = {} if (class_name, meth_name) in annotations["methods"]: annotations_meth = annotations["methods"][(class_name, meth_name)] else: annotations_meth = {} meth_name = fdef.name python_code, attributes, _ = make_new_code_method_from_nodes( class_def, fdef) for attr in attributes: if attr not in annotations_class: raise NotImplementedError( f"self.{attr} used but {attr} not in class annotations") types_attrs = { "self_" + attr: annotations_class[attr] for attr in attributes } types_pythran = {**types_attrs, **annotations_meth} # TODO: locals_types for methods locals_types = None signatures_method = self._make_header_from_fdef_annotations( extast.parse(python_code).body[0], [types_pythran], locals_types) str_self_dot_attributes = ", ".join("self." + attr for attr in attributes) args_func = [arg.id for arg in fdef.args.args[1:]] str_args_func = ", ".join(args_func) defaults = fdef.args.defaults nb_defaults = len(defaults) nb_args = len(fdef.args.args) nb_no_defaults = nb_args - nb_defaults - 1 str_args_value_func = [] ind_default = 0 for ind, arg in enumerate(fdef.args.args[1:]): name = arg.id if ind < nb_no_defaults: str_args_value_func.append(f"{name}") else: default = extast.unparse(defaults[ind_default]).strip() str_args_value_func.append(f"{name}={default}") ind_default += 1 str_args_value_func = ", ".join(str_args_value_func) if str_self_dot_attributes: str_args_backend_func = ", ".join( (str_self_dot_attributes, str_args_func)) else: str_args_backend_func = str_args_func name_var_code_new_method = f"__code_new_method__{class_name}__{meth_name}" self._append_line_header_variable(signatures_method, name_var_code_new_method) python_code += (f'\n{name_var_code_new_method} = """\n\n' f"def new_method(self, {str_args_value_func}):\n" f" return backend_func({str_args_backend_func})" '\n\n"""\n') return signatures_method, format_str(python_code)
def extract_returns_annotation(returns, namespace): if returns is None: return source = extast.unparse(returns) return eval(source, namespace)