Add pixi support

.gitignore

@@ -88,5 +88,6 @@ ENV/
 # Rope project settings
 .ropeproject
 
-# other
+# pixi environments
 .pixi
+*.egg-info

README.md

@@ -420,16 +420,17 @@ See [example](example/) for more information or check the output of `unidep -h`
 <!-- ⚠️ This content is auto-generated by `markdown-code-runner`. -->
 ```bash
 usage: unidep [-h]
-              {merge,install,install-all,conda-lock,pip-compile,pip,conda,version} ...
+              {merge,install,install-all,conda-lock,pixi-lock,pip-compile,pip,conda,version} ...
 
 Unified Conda and Pip requirements management.
 
 positional arguments:
-  {merge,install,install-all,conda-lock,pip-compile,pip,conda,version}
+  {merge,install,install-all,conda-lock,pixi-lock,pip-compile,pip,conda,version}
                         Subcommands
     merge               Combine multiple (or a single) `requirements.yaml` or
                         `pyproject.toml` files into a single Conda installable
-                        `environment.yaml` file.
+                        `environment.yaml` file or Pixi installable
+                        `pixi.toml` file.
     install             Automatically install all dependencies from one or
                         more `requirements.yaml` or `pyproject.toml` files.
                         This command first installs dependencies with Conda,
@@ -448,6 +449,11 @@ positional arguments:
                         lock.yml` files for each `requirements.yaml` or
                         `pyproject.toml` file consistent with the global lock
                         file.
+    pixi-lock           Generate a global `pixi.lock` file for a collection of
+                        `requirements.yaml` or `pyproject.toml` files.
+                        Additionally, create individual `pixi.lock` files for
+                        each `requirements.yaml` or `pyproject.toml` file
+                        consistent with the global lock file.
     pip-compile         Generate a fully pinned `requirements.txt` file from
                         one or more `requirements.yaml` or `pyproject.toml`
                         files using `pip-compile` from `pip-tools`. This
@@ -482,30 +488,33 @@ See `unidep merge -h` for more information:
 <!-- ⚠️ This content is auto-generated by `markdown-code-runner`. -->
 ```bash
 usage: unidep merge [-h] [-o OUTPUT] [-n NAME] [--stdout]
-                    [--selector {sel,comment}] [-d DIRECTORY] [--depth DEPTH]
-                    [-v]
+                    [--selector {sel,comment}] [--pixi] [-d DIRECTORY]
+                    [--depth DEPTH] [-v]
                     [-p {linux-64,linux-aarch64,linux-ppc64le,osx-64,osx-arm64,win-64}]
                     [--skip-dependency SKIP_DEPENDENCY]
                     [--ignore-pin IGNORE_PIN] [--overwrite-pin OVERWRITE_PIN]
 
 Combine multiple (or a single) `requirements.yaml` or `pyproject.toml` files
-into a single Conda installable `environment.yaml` file. Example usage:
-`unidep merge --directory . --depth 1 --output environment.yaml` to search for
-`requirements.yaml` or `pyproject.toml` files in the current directory and its
-subdirectories and create `environment.yaml`. These are the defaults, so you
-can also just run `unidep merge`.
+into a single Conda installable `environment.yaml` file or Pixi installable
+`pixi.toml` file. Example usage: `unidep merge --directory . --depth 1
+--output environment.yaml` to search for `requirements.yaml` or
+`pyproject.toml` files in the current directory and its subdirectories and
+create `environment.yaml`. These are the defaults, so you can also just run
+`unidep merge`.
 
 options:
   -h, --help            show this help message and exit
   -o, --output OUTPUT   Output file for the conda environment, by default
-                        `environment.yaml`
+                        `environment.yaml` or `pixi.toml` if `--pixi` is used
   -n, --name NAME       Name of the conda environment, by default `myenv`
   --stdout              Output to stdout instead of a file
   --selector {sel,comment}
                         The selector to use for the environment markers, if
                         `sel` then `- numpy # [linux]` becomes `sel(linux):
                         numpy`, if `comment` then it remains `- numpy #
                         [linux]`, by default `sel`
+  --pixi                Generate a `pixi.toml` file instead of
+                        `environment.yaml`
   -d, --directory DIRECTORY
                         Base directory to scan for `requirements.yaml` or
                         `pyproject.toml` file(s), by default `.`

example/environment.yaml

@@ -18,7 +18,8 @@ dependencies:
   - pytest
   - pytest-cov
   - pip:
-    - unidep
+    - unidep; sys_platform == 'linux' and platform_machine == 'x86_64'
+    - unidep; sys_platform == 'darwin'
     - markdown-code-runner
     - numthreads
     - yaml2bib; sys_platform == 'linux' and platform_machine == 'x86_64'

example/hatch_project/requirements.yaml

@@ -3,7 +3,7 @@ channels:
   - conda-forge
 dependencies:
   - conda: adaptive-scheduler  # [linux64]
-  - pip: unidep
+  - pip: unidep  # [linux64]
   - numpy >=1.21
   - hpc05  # [linux64]
   - pandas >=1,<3

pixi_create_sub_lock_file.py

@@ -0,0 +1,154 @@
+"""Create a subset of a lock file with a subset of packages."""
+
+from __future__ import annotations
+
+import asyncio
+import json
+import os
+import tempfile
+from collections import defaultdict
+
+from rattler import (
+    Environment,
+    GenericVirtualPackage,
+    LockFile,
+    Platform,
+    Version,
+    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: str) -> dict[str, dict]:
+    """Create repodata from a pixi lock file."""
+    lock_file = LockFile.from_path(lock_file_path)
+    env = lock_file.default_environment()
+    repodata = {}
+    for platform in env.platforms():
+        subdir = str(platform)
+        packages = env.conda_repodata_records_for_platform(platform)
+        if not packages:
+            continue
+
+        repodata[subdir] = {
+            "info": {
+                "subdir": subdir,
+                "base_url": f"https://conda.anaconda.org/conda-forge/{subdir}",
+            },
+            "packages": {
+                f"{pkg.name.normalized}-{pkg.version}-{pkg.build}.conda": {
+                    "build": pkg.build,
+                    "build_number": pkg.build_number,
+                    "depends": pkg.depends,
+                    "constrains": pkg.constrains,
+                    "license": pkg.license,
+                    "license_family": pkg.license_family,
+                    "md5": pkg.md5.hex() if pkg.md5 else None,
+                    "name": pkg.name.normalized,
+                    "sha256": pkg.sha256.hex() if pkg.sha256 else None,
+                    "size": pkg.size,
+                    "subdir": pkg.subdir,
+                    "timestamp": int(pkg.timestamp.timestamp() * 1000)
+                    if pkg.timestamp
+                    else None,
+                    "version": str(pkg.version),
+                }
+                for pkg in packages
+            },
+            "repodata_version": 2,
+        }
+    return repodata
+
+
+def _version_requirement_to_lowest_version(version: str | None) -> str | None:
+    if version is None:
+        return None
+    if version.startswith(">="):
+        version = version[2:]
+    if version.startswith("=="):
+        version = version[2:]
+    version = version.split(",")[0]
+    return version  # noqa: RET504
+
+
+def all_virtual_packages(env: Environment) -> dict[Platform, set[str]]:
+    """Get all virtual packages from an environment."""
+    virtual_packages = defaultdict(set)
+    for platform, packages in env.packages_by_platform().items():
+        for package in packages:
+            if not package.is_conda:
+                continue
+            repo_record = package.as_conda()
+            for dep in repo_record.depends:
+                spec = MatchSpec(dep)
+                if spec.name.normalized.startswith("__"):
+                    version = _version_requirement_to_lowest_version(spec.version)
+                    virtual_package = GenericVirtualPackage(
+                        spec.name,
+                        version=Version(version or "0"),
+                        build_string=spec.build or "*",
+                    )
+                    virtual_packages[platform].add(virtual_package)
+    return virtual_packages
+
+
+async def create_subset_lock_file(
+    original_lock_file_path: str,
+    required_packages: list[str],
+    platform: Platform,
+) -> LockFile:
+    """Create a new lock file with a subset of packages from original lock file."""
+    original_lock_file = LockFile.from_path(original_lock_file_path)
+    env = original_lock_file.default_environment()
+    conda_records = env.conda_repodata_records_for_platform(platform)
+    if conda_records is None:
+        msg = f"No conda records found for platform {platform}"
+        raise ValueError(msg)
+    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)
+
+    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(pkg) for pkg in required_packages]
+    virtual_packages = all_virtual_packages(env)[platform]
+
+    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)  # noqa: PTH108
+    return new_lock_file
+
+
+async def main() -> None:
+    """Example usage of create_subset_lock_file."""
+    original_lock_file_path = "pixi.lock"
+    required_packages = ["tornado", "scipy", "ipykernel", "adaptive"]
+    platform = Platform("linux-64")
+    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())

tests/simple_monorepo/common-requirements.yaml

@@ -5,3 +5,6 @@ channels:
   - conda-forge
 dependencies:
   - conda: python_abi
+platforms:
+  - osx-64
+  - osx-arm64