Create uv managed python shims (similar to pyenv shims)

__uv-python-shims.sh

#!/bin/bash

<<comment
  ============================================================
    Last Edit:    11/21/2024
    Created by:   instance.id (https://github.com/instance.id)
    Platform:     Linux
    Filename:     uv-python-shims
  ============================================================
  .DESCRIPTION
    Creates shims in $HOME/.local/bin for Python versions installed with uv,
    and updates them to reflect the current state of uv installations.

  .EXAMPLE
    uv-python-shims
    uv-python-shims 3.11

  .NOTES
    Inspired by: https://bluesock.org/~willkg/blog/dev/switch_pyenv_to_uv.html
comment

# Directory constants and temporary file setup
LOCALBIN="${HOME}/.local/bin"
UVDIR="$(uv python dir)"
MANAGED_SHIMS=$(mktemp)

# Create a single shim file with appropriate paths and permissions
create_shim() {
  local shim_path="$1"
  local python_path="$2"
  local uv_version_dir="$3"

  cat <<EOF >"${shim_path}"
#!/bin/bash
# Shim generated by uv-python-shims
export PYTHONHOME="${uv_version_dir}"
exec "${python_path}" "\$@"
EOF
  chmod +x "${shim_path}"
  echo "${shim_path}" >>"${MANAGED_SHIMS}"
}

# Create matching config shim if config exists for the Python version
create_config_shim() {
  local config_path="$1"
  local config_shim="$2"
  local uv_version_dir="$3"

  if [[ -f "${config_path}" ]]; then
    create_shim "${config_shim}" "${config_path}" "${uv_version_dir}"
  fi
}

# Create generic python/python3 shims for the desired version
create_default_shims() {
  local python_path="$1"
  local uv_version_dir="$2"

  create_shim "${LOCALBIN}/python3" "${python_path}" "${uv_version_dir}"
  create_shim "${LOCALBIN}/python" "${python_path}" "${uv_version_dir}"
}

# Remove shims for uv Pythons that are no longer installed
cleanup_old_shims() {
  for shim in "${LOCALBIN}"/python*; do
    if [[ -f "${shim}" ]]; then
      if grep -q '# Shim generated by uv-python-shims' "${shim}"; then
        # Check if this shim is still managed
        if ! grep -qx "${shim}" "${MANAGED_SHIMS}"; then
          echo "Removing obsolete shim: ${shim}"
          rm -f "${shim}"
        fi
      fi
    fi
  done
}

# Iterate over installed Python versions and create appropriate shims
find_python_versions() {
  local desired_version="$1"
  local desired_python_path=""

  for python_path in "${UVDIR}"/*/bin/python*; do
    # Skip non-files and python-config files
    if [[ ! -f "${python_path}" ]] || [[ "${python_path}" =~ .*config$ ]]; then
      continue
    fi

    local version_name=$(basename "${python_path}")
    local uv_version_dir=$(dirname "$(dirname "${python_path}")")
    local shim_path="${LOCALBIN}/${version_name}"
    local config_path="${python_path}-config"
    local config_shim="${LOCALBIN}/${version_name}-config"

    # Check if this is the desired version
    if [[ -n "${desired_version}" && "${python_path}" =~ python${desired_version}$ ]]; then
      desired_python_path="${python_path}"
    fi

    create_shim "${shim_path}" "${python_path}" "${uv_version_dir}"
    create_config_shim "${config_path}" "${config_shim}" "${uv_version_dir}"
  done

  # Create default python/python3 shims if desired version found
  if [[ -n "${desired_python_path}" ]]; then
    local uv_version_dir=$(dirname "$(dirname "${desired_python_path}")")
    create_default_shims "${desired_python_path}" "${uv_version_dir}"
  fi
}

# Main entry point
main() {
  # Ensure the local bin directory exists
  mkdir -p "${LOCALBIN}"
  find_python_versions "$1"
  cleanup_old_shims
  # Clean up the temporary file
  rm -f "${MANAGED_SHIMS}"
}

main "$@"

_uv-python-functions.zsh

# ~/.zshrc or ~/.bashrc

# install or uninstall uv python versions using these functions
# to auto update the shims upon completion
function uv-python-install() {
  uv python install "$@"
  $ZDOTDIR/scripts/uv-python-shims
}

function uv-python-uninstall() {
  uv python uninstall "$@"
  $ZDOTDIR/scripts/uv-python-shims
}

# set default python version for bin/python and bin/python3
function uv-python-default() {
  # ensure that the passed python version is installed
  uv-python-install "$@"
  $ZDOTDIR/scripts/uv-python-shims "$@"
}

Refactored and added the ability to set the default python/python3 version.