make the current virtual environment aware of the a rez packages requirements.

rez-venv.py

#! /usr/bin/python3
"""
This script parses a `package.py` file add adds its dependencies to a virtual environment.

- find package.py in the current directory or any parent directory
- extract the 'requires' variable from package.py
- build the environment using `rez env`
- get the site-packages path for venv's Python interpreter
- create a rez.pth file in the site-packages directory
- write the paths from the built environment to rez.pth
"""

import os
import ast
import json
import pathlib
import subprocess


def find_package_file():
    """
    starting at the cwd walk up the directory tree and find package.py
    """
    current_path = pathlib.Path.cwd()
    while current_path != current_path.parent:
        package_file = current_path / "package.py"
        if package_file.exists():
            return package_file
        current_path = current_path.parent
    raise FileNotFoundError("package.py not found in the directory tree.")


def extract_requirements(package_file):
    """
    Extracts requirements from the package.py file.

    - load package.py as text
    - parse it as an AST
    - find the requires attribute, it is a top-level variable
    - evaluate the requires attribute to get the list of requirements
    - return the list of requirements

    """
    with open(package_file, "r") as file:
        package_content = file.read()

    # Parse the content into an AST
    tree = ast.parse(package_content)

    # Find the 'requires' variable in the AST
    for node in ast.walk(tree):
        if isinstance(node, ast.Assign):
            for target in node.targets:
                if isinstance(target, ast.Name) and target.id == "requires":
                    # Evaluate the value of 'requires'
                    return ast.literal_eval(node.value)

    raise ValueError("No 'requires' variable found in package.py.")


def build_environment(packages: list[str]):
    """
    build a rez environment using the given packages and extract the resulting
    python path.
    """
    cmd = ["rez", "env", *packages, "--", "bash", "-c", "echo $PYTHONPATH"]
    result = subprocess.run(cmd, capture_output=True, text=True)

    if result.returncode != 0:
        raise RuntimeError("Failed to build environment")

    return result.stdout.strip().split(os.pathsep)


def get_site_packages_path(venv_path):
    """
    get the site-packages directory for the current virtual environment.
    """    
    python_executable = os.path.join(venv_path, "bin", "python")
    if not os.path.exists(python_executable):
        raise FileNotFoundError(f"Python executable not found at {python_executable}")

    cmd = [python_executable, "-c", "import site; print(site.getsitepackages()[0])"]
    result = subprocess.run(cmd, capture_output=True, text=True)

    if result.returncode != 0:
        raise RuntimeError("Failed to get site-packages path")

    return result.stdout.strip()


def get_venv_path():
    """
    Try to get the path to the virtual environment.
    """
    paths = [
        os.environ.get("VIRTUAL_ENV", ""),
        os.path.join(os.getcwd(), ".venv"),
        os.path.join(os.getcwd(), "venv"),
    ]
    for path in paths:
        if os.path.exists(path):
            return path
    return None


def main():
    venv_path = get_venv_path()
    package_file = find_package_file()
    requirements = extract_requirements(package_file)
    env_paths = build_environment(requirements)

    site_packages_path = get_site_packages_path(venv_path)
    pth_file = os.path.join(site_packages_path, "rez.pth")

    with open(pth_file, "w") as f:
        for path in env_paths:
            f.write(path + "\n")


if __name__ == "__main__":
    try:
        main()
    except Exception as e:
        print(f"Error: {e}")
        raise

    else:
        print("rez.pth file created successfully.")