def __init__(self, backend_name: str, frame=None):
self.backend_name = backend_name
if frame is None:
frame = get_frame(1)
self.filename = _get_filename_from_frame(frame)
if any(
self.filename.startswith(start)
for start in ("<ipython-", "/tmp/ipykernel_")):
self.is_dummy_file = True
self._ipython_src, self.filename = get_info_from_ipython()
self.module_name = self.filename
else:
self.is_dummy_file = False
self.module_name = get_module_name(frame)
modules_backends[backend_name][self.module_name] = self
self.used_functions = {}
self.jit_functions = {}
(
jitted_dicts,
codes_dependance,
codes_dependance_classes,
code_ext,
special,
) = analysis_jit(self.get_source(), self.filename, backend_name)
self.info_analysis = {
"jitted_dicts": jitted_dicts,
"codes_dependance": codes_dependance,
"codes_dependance_classes": codes_dependance_classes,
"special": special,
}
self.backend = backend = backends[backend_name]
path_jit = mpi.Path(backend.jit.path_base)
path_jit_class = mpi.Path(backend.jit.path_class)
# TODO: check if these files have to be written here...
# Write exterior code for functions
for file_name, code in code_ext["function"].items():
path_ext = path_jit / self.module_name.replace(".", os.path.sep)
path_ext_file = path_ext / (file_name + ".py")
write_if_has_to_write(path_ext_file, format_str(code), logger.info)
# Write exterior code for classes
for file_name, code in code_ext["class"].items():
path_ext = path_jit_class / self.module_name.replace(
".", os.path.sep)
path_ext_file = path_ext / (file_name + ".py")
write_if_has_to_write(path_ext_file, format_str(code), logger.info)
def compile_extension(
self,
path_backend,
name_ext_file=None,
native=False,
xsimd=False,
openmp=False,
str_accelerator_flags: Optional[str] = None,
parallel=True,
force=True,
):
if name_ext_file is None:
name_ext_file = self.name_ext_from_path_backend(path_backend)
with open(path_backend) as file:
source = file.read()
source = source.replace("# __protected__ ", "")
with open(path_backend.with_name(name_ext_file), "w") as file:
file.write(format_str(source))
compiling = False
process = None
return compiling, process
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 _make_backend_code(self, path_py, analyse):
"""Create a backend code from a Python file"""
boosted_dicts, code_dependance, annotations, blocks, codes_ext = analyse
boosted_dicts = {
key: value[self.name]
for key, value in boosted_dicts.items()
}
lines_code = ["\n" + code_dependance + "\n"]
lines_header = self._make_first_lines_header()
# Deal with functions
for fdef in boosted_dicts["functions"].values():
signatures_func = self._make_header_1_function(fdef, annotations)
if signatures_func:
lines_header.extend(signatures_func)
code_function = self._make_code_from_fdef_node(fdef)
lines_code.append(code_function)
# Deal with methods
signatures, code_for_meths = self._make_code_methods(
boosted_dicts, annotations, path_py)
lines_code.extend(code_for_meths)
if signatures:
lines_header.extend(signatures)
# Deal with blocks
signatures, code_blocks = self._make_code_blocks(blocks)
lines_code.extend(code_blocks)
if signatures:
lines_header.extend(signatures)
code = "\n".join(lines_code).strip()
if code:
code = self._make_beginning_code() + code
self._append_line_header_variable(lines_header, "__transonic__")
code += f"\n\n__transonic__ = ('{transonic.__version__}',)"
return format_str(
code), codes_ext, "\n".join(lines_header).strip() + "\n"
def make_backend_file(self,
path_py: Path,
analyse=None,
force=False,
log_level=None):
"""Create a Python file from a Python file (if necessary)"""
if log_level is not None:
logger.set_level(log_level)
path_py = Path(path_py)
if not path_py.exists():
raise FileNotFoundError(f"Input file {path_py} not found")
if path_py.absolute().parent.name == f"__{self.name}__":
logger.debug(f"skip file {path_py}")
return None, None, None
if not path_py.name.endswith(".py"):
raise ValueError(
f"transonic only processes Python file. Cannot process {path_py}"
)
path_dir = path_py.parent / str(f"__{self.name}__")
path_backend = (path_dir / path_py.name).with_suffix(
self.suffix_backend)
if not has_to_build(path_backend, path_py) and not force:
logger.warning(f"File {path_backend} already up-to-date.")
return None, None, None
if path_dir is None:
return
if not analyse:
with open(path_py) as file:
code = file.read()
analyse = analyse_aot(code, path_py, self.name)
code_backend, codes_ext, code_header = self._make_backend_code(
path_py, analyse)
if not code_backend:
return
logger.debug(f"code_{self.name}:\n{code_backend}")
for file_name, code in codes_ext["function"].items():
path_ext_file = path_dir / (file_name + ".py")
write_if_has_to_write(path_ext_file, format_str(code), logger.info,
force)
for file_name, code in codes_ext["class"].items():
path_ext_file = (path_dir.parent / f"__{self.name}__" /
(file_name + ".py"))
write_if_has_to_write(path_ext_file, format_str(code), logger.info,
force)
written = write_if_has_to_write(path_backend, code_backend,
logger.info, force)
if not written:
logger.warning(f"Code in file {path_backend} already up-to-date.")
return
if self.suffix_header:
path_header = (path_dir / path_py.name).with_suffix(
self.suffix_header)
write_if_has_to_write(path_header, code_header, logger.info, force)
logger.info(f"File {path_backend} updated")
return path_backend
def _make_code_from_fdef_node(self, fdef):
transformed = TypeHintRemover().visit(fdef)
# convert the AST back to source code
code = extast.unparse(transformed)
return format_str(code)
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 make_code_method_jit(cls, func_name):
func = cls.__dict__[func_name]
func = func.func
new_code, attributes, name_new_func = make_new_code_method_from_objects(
cls, func)
try:
cls_annotations = cls.__annotations__
except AttributeError:
cls_annotations = {}
types_attrs = [
cls_annotations[attr] for attr in attributes if attr in cls_annotations
]
signature = inspect.signature(func)
types_func = [param.annotation
for param in signature.parameters.values()][1:]
types_pythran = types_attrs + types_func
transonic_signatures = "\n"
try:
signatures_as_lists_strings = compute_signatures_from_typeobjects(
types_pythran, base_type_formatter)
except ValueError:
signatures_as_lists_strings = []
for signature_as_strings in signatures_as_lists_strings:
transonic_signatures += (f"# transonic def {name_new_func}(" +
", ".join(signature_as_strings) + ")\n")
new_code = "from transonic import jit\n\n@jit\n" + new_code
python_code = transonic_signatures + "\n" + new_code
str_self_dot_attributes = ", ".join("self." + attr for attr in attributes)
args_func = list(signature.parameters.keys())[1:]
str_args_func = ", ".join(args_func)
str_args_value_func = ""
for param, value in signature.parameters.items():
if param == "self":
continue
elif value.default is value.empty:
str_args_value_func += f"{param}, "
else:
str_args_value_func += f"{param}={value.default}, "
str_args_value_func = str_args_value_func.rstrip(", ")
name_new_method = f"__new_method__{cls.__name__}__{func_name}"
python_code += (
f"\ndef {name_new_method}"
f"(self, {str_args_value_func}):\n"
f" return {name_new_func}({str_self_dot_attributes}, {str_args_func})"
"\n")
python_code = format_str(python_code)
return python_code