/llama-dev/llama_dev/utils.py

import re
import subprocess
import sys
from enum import Enum
from pathlib import Path

import tomli
from packaging import specifiers, version
from packaging.version import Version

DEP_NAME_REGEX = re.compile(r"([^<>=\[\];\s]+)")


class BumpType(str, Enum):
    MAJOR = "major"
    MINOR = "minor"
    PATCH = "patch"


def bump_version(current_version: str, bump_type: BumpType) -> str:
    """Bump a version string according to semver rules."""
    v = Version(current_version)

    # Parse the version components
    release = v.release
    major = release[0] if len(release) > 0 else 0
    minor = release[1] if len(release) > 1 else 0
    micro = release[2] if len(release) > 2 else 0

    version_str = ""
    if bump_type == BumpType.MAJOR:
        version_str = f"{major + 1}.0.0"
    elif bump_type == BumpType.MINOR:
        version_str = f"{major}.{minor + 1}.0"
    elif bump_type == BumpType.PATCH:
        version_str = f"{major}.{minor}.{micro + 1}"

    return version_str


def update_pyproject_version(package_path: Path, new_version: str) -> None:
    """Update the version in a pyproject.toml file."""
    pyproject_path = package_path / "pyproject.toml"

    # Read the file content
    with open(pyproject_path, "r") as f:
        content = f.read()

    pattern = r'^version = "[^"]+"'
    new_content = re.sub(
        pattern, f'version = "{new_version}"', content, flags=re.MULTILINE
    )

    # Write the updated content back
    with open(pyproject_path, "w") as f:
        f.write(new_content)


def package_has_tests(package_path: Path) -> bool:
    """Returns whether a package folder contains a 'tests' subfolder."""
    tests_folder = package_path / "tests"
    return package_path.is_dir() and tests_folder.exists() and tests_folder.is_dir()


def is_llama_index_package(package_path: Path) -> bool:
    """Returns whether a folder contains a 'pyproject.toml' file."""
    pyproject = package_path / "pyproject.toml"
    return package_path.is_dir() and pyproject.exists() and pyproject.is_file()


def load_pyproject(package_path: Path) -> dict:
    """Thin wrapper around tomli.load()."""
    pyproject_path = package_path / "pyproject.toml"
    with open(pyproject_path, "rb") as f:
        return tomli.load(f)


def find_integrations(root_path: Path, recursive=False) -> list[Path]:
    """Find all integrations packages in the repo."""
    package_roots: list[Path] = []
    integrations_root = root_path
    if not recursive:
        integrations_root = integrations_root / "llama-index-integrations"

    for category_path in integrations_root.iterdir():
        if not category_path.is_dir():
            continue

        if category_path.name == "storage":
            # The "storage" category has sub-folders
            package_roots += find_integrations(category_path, recursive=True)
            continue

        for package_name in category_path.iterdir():
            if is_llama_index_package(package_name):
                package_roots.append(package_name)

    return package_roots


def find_utils(root_path: Path) -> list[Path]:
    """Find all llama-index-utils packages in the repo."""
    package_roots: list[Path] = []
    utils_root = root_path / "llama-index-utils"

    for package_name in utils_root.iterdir():
        if is_llama_index_package(package_name):
            package_roots.append(package_name)

    return package_roots


def find_all_packages(root_path: Path) -> list[Path]:
    """Returns a list of all the package folders in the monorepo."""
    return [
        root_path / "llama-index-core",
        *find_integrations(root_path),
        *find_utils(root_path),
        root_path / "llama-index-instrumentation",
    ]


def get_changed_files(repo_root: Path, base_ref: str = "main") -> list[Path]:
    """Use git to get the list of files changed compared to the base branch."""
    try:
        cmd = ["git", "diff", "--name-only", f"{base_ref}...HEAD"]
        result = subprocess.run(cmd, cwd=repo_root, text=True, capture_output=True)
        if result.returncode != 0:
            raise RuntimeError(f"Git command failed: {result.stderr}")

        return [repo_root / Path(f) for f in result.stdout.splitlines() if f.strip()]
    except Exception as e:
        print(f"Exception occurred: {e!s}")
        raise


def get_changed_packages(
    changed_files: list[Path], all_packages: list[Path]
) -> set[Path]:
    """Get the list of package folders containing the path in 'changed_files'."""
    changed_packages: set[Path] = set()

    for file_path in changed_files:
        # Find the package containing this file
        for pkg_dir in all_packages:
            if file_path.absolute().is_relative_to(pkg_dir.absolute()):
                changed_packages.add(pkg_dir)
                break

    return changed_packages


def get_dep_names(pyproject_data: dict) -> set[str]:
    """Load dependencies from pyproject.toml."""
    dependencies: set[str] = set()
    for dep in pyproject_data["project"]["dependencies"]:
        matches = DEP_NAME_REGEX.findall(dep)
        if not matches:
            continue
        dependencies.add(matches[0])
    return dependencies


def is_python_version_compatible(pyproject_data: dict) -> bool:
    """Check if the package is compatible with the current Python version using packaging."""
    # Get current Python version
    current_version = version.Version(
        f"{sys.version_info.major}.{sys.version_info.minor}.{sys.version_info.micro}"
    )

    # Get Python version requirements if they exist
    requires_python = pyproject_data.get("project", {}).get("requires-python")
    if requires_python is None:
        # If no Python version is specified, assume it's compatible
        return True

    try:
        # Parse the version specifier
        spec = specifiers.SpecifierSet(requires_python)

        # Check if the current version satisfies the specifier
        return spec.contains(str(current_version))
    except Exception as e:
        # If there's any error in parsing, log it and assume compatibility
        print(
            f"Warning: Could not parse Python version specifier '{requires_python}': {e}"
        )
        return True


def get_dependants_packages(
    changed_packages: set[Path], all_packages: list[Path]
) -> set[Path]:
    """Get packages containing the files in the changeset."""
    changed_packages_names: set[str] = set()
    for pkg_path in changed_packages:
        pyproject_data = load_pyproject(pkg_path)
        changed_packages_names.add(pyproject_data["project"]["name"])

    dependants_packages: set[Path] = set()
    for pkg_path in all_packages:
        pyproject_data = load_pyproject(pkg_path)
        for dep_name in get_dep_names(pyproject_data):
            if dep_name in changed_packages_names:
                dependants_packages.add(pkg_path)

    return dependants_packages