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)