pixi_create_sub_lock_file.py

import asyncio
import json
import os
from rattler import (
    Environment,
    LockFile,
    Platform,
    VirtualPackage,
    solve_with_sparse_repodata,
)
from rattler.channel import Channel, ChannelConfig
from rattler.match_spec import MatchSpec
from rattler.repo_data import SparseRepoData


def create_repodata_from_pixi_lock(lock_file_path):
    lock_file = LockFile.from_path(lock_file_path)
    env = lock_file.default_environment()
    repodata = {}

    for platform in env.platforms():
        subdir = str(platform)
        repodata[subdir] = {
            "info": {
                "subdir": subdir,
                "base_url": f"https://conda.anaconda.org/conda-forge/{subdir}",
            },
            "packages": {},
        }
        conda_packages = env.conda_repodata_records_for_platform(platform)
        if conda_packages:
            for package in conda_packages:
                filename = (
                    f"{package.name.normalized}-{package.version}-{package.build}.conda"
                )

                repodata[subdir]["packages"][filename] = {
                    "build": package.build,
                    "build_number": package.build_number,
                    "depends": package.depends,
                    "constrains": package.constrains,
                    "license": package.license,
                    "license_family": package.license_family,
                    "md5": package.md5.hex() if package.md5 else None,
                    "name": package.name.normalized,
                    "sha256": package.sha256.hex() if package.sha256 else None,
                    "size": package.size,
                    "subdir": package.subdir,
                    "timestamp": int(package.timestamp.timestamp() * 1000)
                    if package.timestamp
                    else None,
                    "version": str(package.version),
                }

    return repodata


async def create_subset_lock_file(
    original_lock_file_path: str, required_packages: list[str], platform: Platform
):
    original_lock_file = LockFile.from_path(original_lock_file_path)
    default_env = original_lock_file.default_environment()
    conda_records = default_env.conda_repodata_records_for_platform(platform)
    if conda_records is None:
        msg = f"No conda records found for platform {platform}"
        raise ValueError(msg)
    sparse_repo_data = SparseRepoData(
        default_env.channels()[0], platform.subdir, original_lock_file_path
    )
    specs = [MatchSpec(f"{pkg}") for pkg in required_packages]
    solved_records = await solve_with_sparse_repodata(
        specs=specs, sparse_repodata=[sparse_repo_data], locked_packages=conda_records
    )
    new_env = Environment("new_env", {platform: solved_records})
    new_lock_file = LockFile({"new_env": new_env})
    return new_lock_file


async def create_subset_lock_file(
    original_lock_file_path: str, required_packages: list[str], platform: Platform
):
    original_lock_file = LockFile.from_path(original_lock_file_path)
    default_env = original_lock_file.default_environment()
    conda_records = default_env.conda_repodata_records_for_platform(platform)
    if conda_records is None:
        raise ValueError(f"No conda records found for platform {platform}")
    repodata = create_repodata_from_pixi_lock(original_lock_file_path)
    platform_repodata = repodata.get(str(platform))
    if platform_repodata is None:
        msg = f"No repodata found for platform {platform}"
        raise ValueError(msg)

    import tempfile

    with tempfile.NamedTemporaryFile(
        mode="w", delete=False, suffix=".json"
    ) as temp_file:
        json.dump(platform_repodata, temp_file)
        temp_file_path = temp_file.name
    print(f"Temporary repodata file: {temp_file_path}")
    dummy_channel = Channel("dummy", ChannelConfig())
    sparse_repo_data = SparseRepoData(dummy_channel, str(platform), temp_file_path)
    specs = [MatchSpec(f"{pkg}") for pkg in required_packages]
    print(f"Specs: {specs}")
    virtual_packages = VirtualPackage.detect()
    print(f"Detected virtual packages: {virtual_packages}")
    solved_records = await solve_with_sparse_repodata(
        specs=specs,
        sparse_repodata=[sparse_repo_data],
        locked_packages=conda_records,
        virtual_packages=virtual_packages,
    )
    new_env = Environment("new_env", {platform: solved_records})
    new_lock_file = LockFile({"new_env": new_env})
    os.unlink(temp_file_path)
    return new_lock_file


async def main():
    original_lock_file_path = "pixi.lock"
    required_packages = ["pandas"]
    platform = Platform("osx-arm64")
    new_lock_file = await create_subset_lock_file(
        original_lock_file_path,
        required_packages,
        platform,
    )
    new_lock_file.to_path("new_lock_file.lock")


if __name__ == "__main__":
    asyncio.run(main())