def print_dumped(source):
"""Pretty print the AST"""
if isinstance(source, str):
module = extast.parse(source)
if len(module.body) == 1:
node = module.body[0]
else:
node = module
else:
node = source
print(dump(node))
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 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)
code.append(extast.unparse(node).strip())
return "\n".join(code)
if __name__ == "__main__":
code = "a = 1; b = [a, a]\ndef foo():\n return b"
# code = "a = 1; b = len([a, a])\ndef foo():\n return b"
# code = "import numpy as np\na = np.int(1)\ndef foo():\n return np.zeros(a)"
code = """
a = 1
def fooo():
return 1
def foo():
return a + fooo()
def bar():
return foo()
"""
module = extast.parse(code)
function = module.body[3]
capturex = CaptureX((function, ), module)
print(capturex.make_code_external())
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 get_exterior_code(
codes_dependance: dict,
pathfile: str,
previous_file_name=None,
classes: str = None,
relative: bool = None,
jitted_dicts: dict = None,
):
"""Get all imported functions needed by boosted functions and methods at multiple levels,
(i.e get functions needed by functions imported by boosted function) and add them into code_ext
"""
special = []
treated = []
for func, dep in codes_dependance.items():
if not dep:
continue
module_ext = extast.parse(dep)
for node in module_ext.body:
if not isinstance(node, (ast.ImportFrom, ast.Import)):
continue
# get the path of the imported module
file_name, file_path = find_path(node, pathfile)
# a jitted function or method needs another jitted function
if file_name == "transonic":
(
codes_dependance[func],
jitted_dicts,
spe,
treat,
) = adapt_code_dependance(func, codes_dependance[func],
jitted_dicts)
# the "special" list signals that jitted functions has to be written
# from jitted_dicts (see transonic/justintime.py:287)
special = special + spe
treated = treated + treat
for func, dep in codes_dependance.items():
if not dep:
continue
module_ext = extast.parse(dep)
for node in module_ext.body:
if not isinstance(node, (ast.ImportFrom, ast.Import)):
continue
# get the path of the imported module
file_name, file_path = find_path(node, pathfile)
# a jitted function or method needs another jitted function
if not (file_name and file_name not in treated):
continue
new_file_name = f"__ext__{func}__{file_name}"
# get the content of the file
try:
with open(str(file_path), "r") as file:
content = file.read()
except:
raise NotImplementedError(file_name + " can not be found")
mod = extast.parse(content)
# filter the code and add it to code_ext dict
code_ext[classes][new_file_name] = str(
filter_external_code(mod, node.names))
# change imported module names
codes_dependance[func] = change_import_name(
codes_dependance[func], node, func, relative)
# recursively get the exterior codes
if code_ext[classes][new_file_name]:
get_exterior_code(
{func: code_ext[classes][new_file_name]},
pathfile,
new_file_name,
classes,
)
if previous_file_name:
code_ext[classes][previous_file_name] = change_import_name(
code_ext[classes][previous_file_name],
node,
func,
relative,
)
return codes_dependance, code_ext, jitted_dicts, special
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 analysis_jit(code, pathfile, backend_name):
"""Gather the informations for ``@jit`` with an ast analysis"""
debug = logger.debug
debug("extast.parse")
module = extast.parse(code)
debug("compute ancestors and chains")
ancestors, duc, udc = compute_ancestors_chains(module)
jitted_dicts = get_decorated_dicts(module,
ancestors,
duc,
pathfile,
backend_name,
decorator="jit")
jitted_dicts = dict(
functions=jitted_dicts["functions"][backend_name],
functions_ext=jitted_dicts["functions_ext"][backend_name],
methods=jitted_dicts["methods"][backend_name],
classes=jitted_dicts["classes"][backend_name],
)
debug("compute code dependance")
def_nodes_dict = {
key: def_node
for kind in ["functions"]
for key, def_node in jitted_dicts[kind].items()
}
codes_dependance = {}
# get code dependance of a decorated function from a imported module
for func_name, (func_node,
imported_module) in jitted_dicts["functions_ext"].items():
capturex = CaptureX((func_node, ),
imported_module,
consider_annotations=False)
codes_dependance[func_name] = capturex.make_code_external()
# remove the decorator (jit) to compute the code dependance
for key, def_node in def_nodes_dict.items():
def_node.decorator_list = []
capturex = CaptureX(
(def_node, ),
module,
ancestors=ancestors,
defuse_chains=duc,
usedef_chains=udc,
consider_annotations=False,
)
codes_dependance[key] = capturex.make_code_external()
debug(codes_dependance)
def_nodes_dict = {}
for (class_name, method_name), def_node in jitted_dicts["methods"].items():
if class_name not in def_nodes_dict:
def_nodes_dict[class_name] = []
def_nodes_dict[class_name].append(def_node)
codes_dependance_classes = {}
for key, def_nodes in def_nodes_dict.items():
for def_node in def_nodes:
def_node.decorator_list = []
capturex = CaptureX(
def_nodes,
module,
ancestors=ancestors,
defuse_chains=duc,
usedef_chains=udc,
consider_annotations=False,
)
codes_dependance_classes[key] = capturex.make_code_external()
special = []
spe = []
if jitted_dicts["methods"]:
codes_dependance_classes, code_ext, jitted_dicts, spe = get_exterior_code(
codes_dependance_classes,
pathfile,
previous_file_name=None,
classes="class",
relative=False,
jitted_dicts=jitted_dicts,
)
special = special + spe
codes_dependance, code_ext, jitted_dicts, spe = get_exterior_code(
codes_dependance,
pathfile,
previous_file_name=None,
classes="function",
relative=False,
jitted_dicts=jitted_dicts,
)
special = special + spe
return (
jitted_dicts,
codes_dependance,
codes_dependance_classes,
code_ext,
special,
)
