fix conda/mamba env with explicit versions

fix_env.sh

#!/usr/bin/env zsh

# This script fixes the environment.yml file for a conda environment.
# Dr Ting Sun (ting.sun@ucl.ac.uk) + GitHub Copilot
# 09 Nov 2023: Initial version
# 11 Feb 2024:
#   - Better handling of the pip section
#   - Add "-f" flag to force overwrite the original env.yml file


# It reads the package names from the original environment.yml file, fetches the installed versions from the environment,
# and writes the installed versions to a new environment.yml file.
# The pip section from the original environment.yml file is also appended to the new environment.yml file.
# The path to the environment.yml file is passed as an argument to the script.

# Source the .zshrc file to load the mamba function
source ~/.zshrc

# add some help flag to this cmd tool
# if followed by -h or --help, print the help message
if [[ $1 == "-h" || $1 == "--help" ]]; then
    # Display help information about the script
    echo "This script secures dependency versions in the $(environment.yml) file for conda/mamba environments."
    echo "Usage: fix-env [-f] env.yml"
    echo "Options:"
    echo "  -f    Force overwrite the original env.yml file"
    exit 0
fi

if [ $# -eq 0 ]; then
    echo "No arguments supplied. Please provide the path to env.yml file."
    exit 1
fi

# This script extracts the name of the environment from a file.
# It uses the 'awk' command to search for the line containing 'name:' and prints the second field.
# The file path is passed as an argument to the script.
env_name=$(awk '/name:/{print $2}' $1)
echo "Environment name: $env_name"

# Read package names from the original env.yml file
# only include those after "dependencies:" and start with "  - "
contents=$(awk '/dependencies:/{flag=1} flag' $1)
contents=$(echo "$contents" | grep -E "  - [a-zA-Z0-9]")

# The following removes leading hyphens, comments, version numbers, and trailing colons from the contents of a variable named "contents".
# The modified contents are stored in a variable named "packages".
# The sed command is used multiple times to perform the necessary substitutions.

# Removes leading hyphens and spaces from each line of the input contents.
packages=$(echo "$contents" | sed 's/^[[:space:]]*-[[:space:]]*//')

# Remove comments from the "packages" variable
packages=$(echo "$packages" | sed 's/[[:space:]]*#.*$//')

# Removes the trailing equal sign and everything after it from each package in the variable $packages.
packages=$(echo "$packages" | sed 's/[=].*$//')

# Removes everything after the first colon (":") in the variable "packages".
packages=$(echo "$packages" | sed 's/[:].*$//')

# Removes everything after dash ("-") in the variable "packages".
packages=$(echo "$packages" | sed 's/[-][[:space:]]*$//')

# packages=$(echo "$contents" | sed 's/^[[:space:]]*-[[:space:]]*//; s/[[:space:]]*#.*$//' | grep -E "^[a-zA-Z0-9]")
# show the packages
echo "Packages:"
echo $packages
# separate the packages with "|" for the grep command
packages=$(echo "$packages" | tr '\n' '|')

# packages=$(awk -F= '/ - /{print $1}' $1 | sed 's/^[[:space:]]*-[[:space:]]*//; s/[[:space:]]*#.*$//' | tr '\n' '|')
packages=${packages::-1} # Remove the trailing '|'

echo "Packages to be included:"
echo $packages

# Fetch installed versions from the environment

installed_versions=$(mamba list -n $env_name | grep -w -E "\b($packages) ")

echo "Installed versions:"
echo $installed_versions

# Write to a new env.yml file
# if the file exists, delete it
if [ -f fixed_env.yml ]; then
    rm fixed_env.yml
fi

# if the file does not exist, create it
touch fixed_env.yml

# write the first comment line to the new env.yml file
echo "# This file was generated automatically from '$1' " >fixed_env.yml

# Read the env.yml file and write only the first line to the new env.yml file
awk '/name:/' $1 >>fixed_env.yml

echo "channels:" >>fixed_env.yml
echo "  - conda-forge" >>fixed_env.yml
echo "  - defaults" >>fixed_env.yml
echo "dependencies:" >>fixed_env.yml

# initialize a text block to store the pip section
pip_section=""

# Write the installed versions to the new env.yml file
while read -r line; do
    package=$(echo "$line" | awk '{print $1}')
    version=$(echo $line | awk '{print $2}')
    channel=$(echo $line | awk '{print $4}')

    if [[ $channel != pypi* ]]; then
        echo "  - $package=$version" >>fixed_env.yml
    else
        pip_section+="    - $package==$version\n"
    fi
done <<<"$installed_versions"

# Append the pip section from the installed environment to the new env.yml file if it is not empty
if [[ -n $pip_section ]]; then
    echo "  - pip:" >>fixed_env.yml
    echo -e $pip_section >>fixed_env.yml
fi

# if received no error, and with the `-f` flag, remove the original env.yml file and rename the new one to env.yml
if [ $? -eq 0 ]; then
    if [[ $2 == "-f" ]]; then
        rm $1
        mv fixed_env.yml $1
    fi
fi