first commit

utils/modular_integrations.py

@@ -0,0 +1,184 @@
+import os
+
+import libcst as cst
+
+
+# Files from external libraries that should not be tracked
+# E.g. for habana, we don't want to track the dependencies from `modeling_all_models.py` as it is not part of the transformers library
+EXCLUDED_EXTERNAL_FILES = {
+    "habana": [{"name": "modeling_all_models", "type": "modeling"}],
+}
+
+
+def convert_relative_import_to_absolute(
+    import_node: cst.ImportFrom,
+    file_path: str,
+    package_name: str | None = "transformers",
+) -> cst.ImportFrom:
+    """
+    Convert a relative libcst.ImportFrom node into an absolute one,
+    using the file path and package name.
+
+    Args:
+        import_node: A relative import node (e.g. `from ..utils import helper`)
+        file_path: Path to the file containing the import (can be absolute or relative)
+        package_name: The top-level package name (e.g. 'myproject')
+
+    Returns:
+        A new ImportFrom node with the absolute import path
+    """
+    if not (import_node.relative and len(import_node.relative) > 0):
+        return import_node  # Already absolute
+
+    file_path = os.path.abspath(file_path)
+    rel_level = len(import_node.relative)
+
+    # Strip file extension and split into parts
+    file_path_no_ext = file_path.removesuffix(".py")
+    file_parts = file_path_no_ext.split(os.path.sep)
+
+    # Ensure the file path includes the package name
+    if package_name not in file_parts:
+        raise ValueError(f"Package name '{package_name}' not found in file path '{file_path}'")
+
+    # Slice file_parts starting from the package name
+    pkg_index = file_parts.index(package_name)
+    module_parts = file_parts[pkg_index + 1 :]  # e.g. ['module', 'submodule', 'foo']
+    if len(module_parts) < rel_level:
+        raise ValueError(f"Relative import level ({rel_level}) goes beyond package root.")
+
+    base_parts = module_parts[:-rel_level]
+
+    # Flatten the module being imported (if any)
+    def flatten_module(module: cst.BaseExpression | None) -> list[str]:
+        if not module:
+            return []
+        if isinstance(module, cst.Name):
+            return [module.value]
+        elif isinstance(module, cst.Attribute):
+            parts = []
+            while isinstance(module, cst.Attribute):
+                parts.insert(0, module.attr.value)
+                module = module.value
+            if isinstance(module, cst.Name):
+                parts.insert(0, module.value)
+            return parts
+        return []
+
+    import_parts = flatten_module(import_node.module)
+
+    # Combine to get the full absolute import path
+    full_parts = [package_name] + base_parts + import_parts
+
+    # Handle special case where the import comes from a namespace package (e.g. optimum with `optimum.habana`, `optimum.intel` instead of `src.optimum`)
+    if package_name != "transformers" and file_parts[pkg_index - 1] != "src":
+        full_parts = [file_parts[pkg_index - 1]] + full_parts
+
+    # Build the dotted module path
+    dotted_module: cst.BaseExpression | None = None
+    for part in full_parts:
+        name = cst.Name(part)
+        dotted_module = name if dotted_module is None else cst.Attribute(value=dotted_module, attr=name)
+
+    # Return a new ImportFrom node with absolute import
+    return import_node.with_changes(module=dotted_module, relative=[])
+
+
+def convert_to_relative_import(import_node: cst.ImportFrom, file_path: str, package_name: str) -> cst.ImportFrom:
+    """
+    Convert an absolute import to a relative one if it belongs to `package_name`.
+
+    Parameters:
+    - node: The ImportFrom node to possibly transform.
+    - file_path: Absolute path to the file containing the import (e.g., '/path/to/mypackage/foo/bar.py').
+    - package_name: The top-level package name (e.g., 'mypackage').
+
+    Returns:
+    - A possibly modified ImportFrom node.
+    """
+    if import_node.relative:
+        return import_node  # Already relative import
+
+    # Extract module name string from ImportFrom
+    def get_module_name(module):
+        if isinstance(module, cst.Name):
+            return module.value, [module.value]
+        elif isinstance(module, cst.Attribute):
+            parts = []
+            while isinstance(module, cst.Attribute):
+                parts.append(module.attr.value)
+                module = module.value
+            if isinstance(module, cst.Name):
+                parts.append(module.value)
+            parts.reverse()
+            return ".".join(parts), parts
+        return "", None
+
+    module_name, submodule_list = get_module_name(import_node.module)
+
+    # Check if it's from the target package
+    if (
+        not (module_name.startswith(package_name + ".") or module_name.startswith("optimum." + package_name + "."))
+        and module_name != package_name
+    ):
+        return import_node  # Not from target package
+
+    # Locate the package root inside the file path
+    norm_file_path = os.path.normpath(file_path)
+    parts = norm_file_path.split(os.sep)
+
+    try:
+        pkg_index = parts.index(package_name)
+    except ValueError:
+        # Package name not found in path — assume we can't resolve relative depth
+        return import_node
+
+    # Depth is how many directories after the package name before the current file
+    depth = len(parts) - pkg_index - 1  # exclude the .py file itself
+    for i, submodule in enumerate(parts[pkg_index + 1 :]):
+        if submodule == submodule_list[2 + i]:
+            depth -= 1
+        else:
+            break
+
+    # Create the correct number of dots
+    relative = [cst.Dot()] * depth if depth > 0 else [cst.Dot()]
+
+    # Strip package prefix from import module path
+    if module_name.startswith("optimum." + package_name + "."):
+        stripped_name = module_name[len("optimum." + package_name) :].lstrip(".")
+    else:
+        stripped_name = module_name[len(package_name) :].lstrip(".")
+
+    # Build new module node
+    if stripped_name == "":
+        new_module = None
+    else:
+        name_parts = stripped_name.split(".")[i:]
+        new_module = cst.Name(name_parts[0])
+        for part in name_parts[1:]:
+            new_module = cst.Attribute(value=new_module, attr=cst.Name(part))
+
+    return import_node.with_changes(module=new_module, relative=relative)
+
+
+class AbsoluteImportTransformer(cst.CSTTransformer):
+    def __init__(self, relative_path: str, source_library: str):
+        super().__init__()
+        self.relative_path = relative_path
+        self.source_library = source_library
+
+    def leave_ImportFrom(self, original_node: cst.ImportFrom, updated_node: cst.ImportFrom) -> cst.ImportFrom:
+        return convert_relative_import_to_absolute(
+            import_node=updated_node, file_path=self.relative_path, package_name=self.source_library
+        )
+
+
+class RelativeImportTransformer(cst.CSTTransformer):
+    def __init__(self, relative_path: str, source_library: str):
+        super().__init__()
+        self.relative_path = relative_path
+        self.source_library = source_library
+
+    def leave_ImportFrom(self, original_node: cst.ImportFrom, updated_node: cst.ImportFrom) -> cst.ImportFrom:
+        return convert_to_relative_import(updated_node, self.relative_path, self.source_library)