nix shell runner script for python and packages and optional venv

pysh

#!/usr/bin/env python
"""
pysh

This script provides a command-line interface (CLI) for creating and running a 
nix-shell environment with specified Python packages. It supports using different 
Python versions and can optionally create a temporary virtual environment (venv) 
to install packages.

Purpose:
--------
The purpose of this script is to simplify the setup and execution of Python 
development environments using nix-shell. Users can specify the Python version 
and a list of packages to include in the environment. An optional feature allows 
for the creation of a temporary virtual environment to further isolate package 
installations.

Usage:
------
1. Basic usage with default Python version (3.11) and specified packages:
    ```sh
    ./pysh package1 package2
    ```

2. Specify a different Python version:
    ```sh
    ./pysh --python 312 package1 package2
    ```

3. Create a temporary virtual environment:
    ```sh
    ./pysh --use-venv package1 package2
    ```

Arguments:
----------
- `-p`, `--python`: Specify the Python version (default: 311).
- `-v`, `--use-venv`: Create a temporary virtual environment.
- `packages`: List of additional packages to include in the nix-shell environment.

Functions:
----------
- `build_nix_command(python_version: str, packages: List[str], use_venv: bool = False, venv_dir: str | None = None) -> str`:
    Constructs the nix-shell command based on the provided arguments.

- `main() -> None`:
    Parses command-line arguments, constructs the nix-shell command, and runs it.

Example:
--------
To run nix-shell with Python 3.11 and install requests and numpy:
    ```sh
    ./pysh requests numpy
    ```

To run nix-shell with Python 3.12 and install requests in a temporary virtual environment:
    ```sh
    ./pysh --python 312 --use-venv requests
    ```

Dependencies:
-------------
- `tempfile`: For creating a temporary directory.
- `argparse`: For parsing command-line arguments.
- `subprocess`: For running shell commands.
- `sys`: For system-specific parameters and functions.
"""

import tempfile

import argparse
import subprocess
import sys
from typing import List


def build_nix_command(
    python_version: str,
    packages: List[str],
    use_venv: bool = False,
    venv_dir: str | None = None,
) -> str:
    """
    Build the nix-shell command.

    Args:
        python_version (str): The version of Python to use.
        packages (List[str]): List of packages to include.
        use_venv (bool): Whether to create temporary venv and install packages inside it

    Returns:
        str: The nix-shell command.
    """
    base_package = f"python{python_version}Packages.pip"
    if use_venv:
        if venv_dir is None:
            raise ValueError("venv_dir cannot be None if use_venv is True")
        cmds: List[str] = [
            f"python -m venv {venv_dir}",
            f"source {venv_dir}/bin/activate",
        ]

        if packages:
            cmds.append("pip install " + " ".join(packages))

        cmds.append("return")
        command = f"nix-shell -p {base_package} --command '{';'.join(cmds)}'"
    else:
        packages_cmd = " ".join(
            f"python{python_version}Packages.{pkg}" for pkg in packages
        )
        command = f"nix-shell -p {base_package} {packages_cmd}"
    return command


def main() -> None:
    parser = argparse.ArgumentParser(
        description="pysh CLI for nix-shell with Python packages"
    )
    _ = parser.add_argument(
        "-p",
        "--python",
        type=str,
        default="311",
        help="Python version (e.g., 311, 312)",
    )
    _ = parser.add_argument(
        "-v",
        "--use-venv",
        default=False,
        action="store_true",
        help="Create temporary virtualenv",
    )
    parser.add_argument("packages", nargs="*", help="Additional packages to include")

    args = parser.parse_args()

    with tempfile.TemporaryDirectory() as venv_dir:
        command = build_nix_command(
            args.python, args.packages, args.use_venv, venv_dir=venv_dir
        )
        console_output = f"Running command: {command}"
        print(console_output)

        try:
            subprocess.run(command, shell=True, check=True)
        except subprocess.CalledProcessError as e:
            print(f"Error: {e}", file=sys.stderr)
            sys.exit(e.returncode)


if __name__ == "__main__":
    main()