[WIP] Fix #708

njit_fastmath.py

@@ -0,0 +1,72 @@
+import pathlib
+
+from utils import check_callees, check_functions
+
+stumpy_path = pathlib.Path(__file__).parent / "stumpy"
+filepaths = sorted(f for f in pathlib.Path(stumpy_path).iterdir() if f.is_file())
+
+all_functions = {}
+
+ignore = ["__init__.py", "__pycache__"]
+for filepath in filepaths:
+    file_name = filepath.name
+    if file_name not in ignore and str(filepath).endswith(".py"):
+        prefix = file_name.replace(".py", "")
+
+        func_names, is_njit, fastmath_values = check_functions(filepath)
+        func_names = [f"{prefix}.{fn}" for fn in func_names]
+
+        all_functions[file_name] = {
+            "func_names": func_names,
+            "is_njit": is_njit,
+            "fastmath_values": fastmath_values,
+        }
+
+all_stumpy_functions = set()
+for file_name, file_functions_metadata in all_functions.items():
+    all_stumpy_functions.update(file_functions_metadata["func_names"])
+
+all_stumpy_functions = list(all_stumpy_functions)
+all_stumpy_functions_no_prefix = [f.split(".")[-1] for f in all_stumpy_functions]
+
+
+# output 1: func_metadata
+func_metadata = {}
+for file_name, file_functions_metadata in all_functions.items():
+    for i, f in enumerate(file_functions_metadata["func_names"]):
+        is_njit = file_functions_metadata["is_njit"][i]
+        fastmath_value = file_functions_metadata["fastmath_values"][i]
+        func_metadata[f] = [is_njit, fastmath_value]
+
+
+# output 2: func_callers
+func_callers = {}
+for f in func_metadata.keys():
+    func_callers[f] = []
+
+for filepath in filepaths:
+    file_name = filepath.name
+    if file_name in ignore or not str(filepath).endswith(".py"):
+        continue
+
+    prefix = file_name.replace(".py", "")
+    callees = check_callees(filepath)
+
+    current_callers = set(callees.keys())
+    for caller, callee_set in callees.items():
+        s = list(callee_set.intersection(all_stumpy_functions_no_prefix))
+        if len(s) == 0:
+            continue
+
+        for c in s:
+            if c in current_callers:
+                c_name = prefix + "." + c
+            else:
+                idx = all_stumpy_functions_no_prefix.index(c)
+                c_name = all_stumpy_functions[idx]
+
+            func_callers[c_name].append(f"{prefix}.{caller}")
+
+
+for f, callers in func_callers.items():
+    func_callers[f] = list(set(callers))

utils.py

@@ -0,0 +1,146 @@
+import ast
+
+
+def check_fastmath(decorator):
+    """
+    For the given `decorator` node with type `ast.Call`,
+    return the value of the `fastmath` argument if it exists.
+    Otherwise, return `None`.
+    """
+    fastmath_value = None
+    for n in ast.iter_child_nodes(decorator):
+        if isinstance(n, ast.keyword) and n.arg == "fastmath":
+            if isinstance(n.value, ast.Constant):
+                fastmath_value = n.value.value
+            elif isinstance(n.value, ast.Set):
+                fastmath_value = set(item.value for item in n.value.elts)
+            else:
+                pass
+            break
+
+    return fastmath_value
+
+
+def check_njit(fd):
+    """
+    For the given `fd` node with type `ast.FunctionDef`,
+    return the node of the `njit` decorator if it exists.
+    Otherwise, return `None`.
+    """
+    decorator_node = None
+    for decorator in fd.decorator_list:
+        if not isinstance(decorator, ast.Call):
+            continue
+
+        obj = decorator.func
+        if isinstance(obj, ast.Attribute):
+            name = obj.attr
+        elif isinstance(obj, ast.Subscript):
+            name = obj.value.id
+        elif isinstance(obj, ast.Name):
+            name = obj.id
+        else:
+            msg = f"The type {type(obj)} is not supported."
+            raise ValueError(msg)
+
+        if name == "njit":
+            decorator_node = decorator
+            break
+
+    return decorator_node
+
+
+def check_functions(filepath):
+    """
+    For the given `filepath`, return the function names,
+    whether the function is decorated with `@njit`,
+    and the value of the `fastmath` argument if it exists
+
+    Parameters
+    ----------
+    filepath : str
+        The path to the file
+
+    Returns
+    -------
+    func_names : list
+        List of function names
+
+    is_njit : list
+        List of boolean values indicating whether the function is decorated with `@njit`
+
+    fastmath_value : list
+        List of values of the `fastmath` argument if it exists
+    """
+    file_contents = ""
+    with open(filepath, encoding="utf8") as f:
+        file_contents = f.read()
+    module = ast.parse(file_contents)
+
+    function_definitions = [
+        node for node in module.body if isinstance(node, ast.FunctionDef)
+    ]
+
+    func_names = [fd.name for fd in function_definitions]
+
+    njit_nodes = [check_njit(fd) for fd in function_definitions]
+    is_njit = [node is not None for node in njit_nodes]
+
+    fastmath_values = [None] * len(njit_nodes)
+    for i, node in enumerate(njit_nodes):
+        if node is not None:
+            fastmath_values[i] = check_fastmath(node)
+
+    return func_names, is_njit, fastmath_values
+
+
+def _get_callees(node, all_functions):
+    for n in ast.iter_child_nodes(node):
+        if isinstance(n, ast.Call):
+            obj = n.func
+            if isinstance(obj, ast.Attribute):
+                name = obj.attr
+            elif isinstance(obj, ast.Subscript):
+                name = obj.value.id
+            elif isinstance(obj, ast.Name):
+                name = obj.id
+            else:
+                msg = f"The type {type(obj)} is not supported"
+                raise ValueError(msg)
+
+            all_functions.append(name)
+
+        _get_callees(n, all_functions)
+
+
+def get_all_callees(fd):
+    """
+    For a given node of type ast.FunctionDef, visit all of its child nodes,
+    and return a list of all of its callees
+    """
+    all_functions = []
+    _get_callees(fd, all_functions)
+
+    return all_functions
+
+
+def check_callees(filepath):
+    """
+    For the given `filepath`, return a dictionary with the key
+    being the function name and the value being a set of function names
+    that are called by the function
+    """
+    file_contents = ""
+    with open(filepath, encoding="utf8") as f:
+        file_contents = f.read()
+    module = ast.parse(file_contents)
+
+    function_definitions = [
+        node for node in module.body if isinstance(node, ast.FunctionDef)
+    ]
+
+    callees = {}
+    for fd in function_definitions:
+        callees[fd.name] = set(get_all_callees(fd))
+
+    return callees