alienv

#!/bin/bash -e

# Load and test Modulefiles created by recipes.
PROG=$(basename "$0")

# If we are on a TTY, enable colors
if test -t 1; then
  EM=`printf '\033[35m'`  # magenta
  ET=`printf '\033[36m'`  # cyan
  ER=`printf '\033[31m'`  # red
  EY=`printf '\033[33m'`  # yellow
  EZ=`printf '\033[m'`    # reset
else
  EM=
  ET=
  ER=
  EY=
  EZ=
fi

# My operating system
UNAME=$(uname)

function printHelp() {
cat >&2 <<EoF
$PROG -- load environment for aliBuild packages through modulefiles

Usage: $0 \\
         [--architecture|-a ${ET}ARCHITECTURE${EZ}] \\
         [--work-dir|-w ${ET}WORKDIR${EZ}]          \\
         [--no-refresh]                   \\
         ${ET}COMMAND...${EZ}

  ${EM}--no-refresh${EZ} skips refreshing the modules directory.

  ${ET}WORKDIR${EZ} defaults to sw.
  ${ET}ARCHITECTURE${EZ} is automatically detected in most cases.

  ${ET}COMMAND...${EZ} might be:

    ${EM}help${EZ}
      This help screen.

    ${EM}enter${EZ} [${ET}--shellrc${EZ}] MODULE1[,MODULE2...]
      Enters a new shell with the given modules loaded.
      Return to the clean environment by exiting the shell with ${ET}exit${EZ}.
      Inside the environment you can use the native ${ET}modulecmd${EZ}.
      By default you enter the same shell type you are in. Override with environment variable ${ET}MODULES_SHELL${EZ}.
      The new shell will not load your shell startup file by default (e.g. ~/.bashrc) to avoid environment conflicts.
      If you want to retain your shell startup files use ${ET}--shellrc${EZ}.
      Alternatively you can use the ${EM}load${EZ} command in the current shell.

    ${EM}setenv${EZ} MODULE1[,MODULE2...] ${ET}-c${EZ} cmdInEnvironment [PARAM1 [PARAM2...]]
      Executes the given command with environment defined by the given modules.
      Everything after ${ET}-c${EZ} is executed as-is.
      Exit code is preserved.
      Example: $ET$0 setenv AliRoot/v5-08-02-1 -c aliroot -b$EZ

    ${EM}printenv${EZ} or ${EM}load${EZ} [${ET}-q${EZ}] MODULE1[,MODULE2...]
      Prints the environment in the current shell for the given modules (${ET}-q${EZ} for quiet).
      This command does not set any environment and it must be executed through ${ET}eval${EZ} to be effective.
      Override shell with the environment variable ${ET}MODULES_SHELL${EZ}.
      Example: ${ET}eval \`$0 load AliRoot/latest\`$EZ (those are backquotes!)

    ${EM}unload${EZ} [${ET}-q${EZ}] MODULE1[,MODULE2...]
      Prints the environment in the current shell for unloading the given modules (${ET}-q${EZ} for quiet).
      This command does not set any environment and it must be executed through ${ET}eval${EZ} to be effective.
      Override shell with the environment variable ${ET}MODULES_SHELL${EZ}.
      Example: ${ET}eval \`$0 unload AliRoot\`$EZ (version can be omitted)

    ${EM}q${EZ} or ${EM}query${EZ} [REGEXP]
      List all available modules, or the ones matching ${ET}REGEXP${EZ} if provided.

    ${EM}list${EZ}
      List loaded modules.

    ${EM}modulecmd${EZ} [PARAM1 [PARAM2...]]]
      Pass all arguments as-is to the ${ET}modulecmd${EZ} command.
      Example: print AliRoot env for zsh: $ET$0 modulecmd zsh load AliRoot/v5-08-02-1$EZ
      Consult ${ET}man modulecmd${EZ} for more information.

    ${EM}shell-helper${EZ}
      Returns a script to be ${ET}eval${EZ}d in your shell rc file to allow easier loading/unloading.
      In your shell rc put:
        ${ET}ALIBUILD_WORK_DIR=<path_to_alibuild_sw_dir>${EZ}
        ${ET}eval "\`alienv shell-helper\`"${EZ}

EoF
[[ -z "$1" ]] || printf "${ER}ERROR: $*${EZ}\n" >&2
}

function installHint() {
  if [[ $UNAME == Darwin ]]; then
    CMD='brew install modules'
  elif which apt-get > /dev/null 2>&1; then
    CMD='apt-get install environment-modules'
  elif which yum > /dev/null 2>&1; then
    CMD='yum install environment-modules'
  fi
  printf "${ER}ERROR: Environment Modules was not found on your system.\n" >&2
  if [[ -z "$CMD" ]]; then
    printf "       The package is usually called ${EM}environment-modules${ER}.\n" >&2
  else
    printf "       Get it with: ${EM}${CMD}\n" >&2
  fi
  printf "${EZ}"
}

function detectShell() {
  # Detect parent shell (fall back to bash)
  [[ -z "$MODULES_SHELL" ]] && MODULES_SHELL=$(ps -e -o pid,command | grep -E "^\s*$PPID\s+" | awk '{print $2}' | sed -e 's/^-\{0,1\}\(.*\)$/\1/')
  case "$MODULES_SHELL" in
    sh)                                      ;;
    csh|tcsh) SHELL_NORC_PARAM=-f            ;;
    ksh)      SHELL_NORC_ENV="ENV=/dev/null" ;;
    zsh)      SHELL_NORC_PARAM=--no-rcs      ;;
    *)        MODULES_SHELL=bash
              SHELL_NORC_PARAM=--norc        ;;
  esac
}

function collectModules() {
  if [[ -z "$NO_REFRESH" ]]; then
    # Collect all modulefiles in one place
    rm -rf $WORK_DIR/MODULES/$ARCHITECTURE
    mkdir -p $WORK_DIR/MODULES/$ARCHITECTURE/BASE
    cat > $WORK_DIR/MODULES/$ARCHITECTURE/BASE/1.0 <<EOF
#%Module1.0
set base_path $WORK_DIR/$ARCHITECTURE
setenv BASEDIR \$base_path
set osname [uname sysname]
set osarchitecture [uname machine]
EOF
    while read PKG; do
      PKGVER=${PKG##*/}
      PKGNAME=${PKG%/*}
      PKGNAME=${PKGNAME##*/}
      [[ ! -e "$PKG/etc/modulefiles/$PKGNAME" ]] && continue
      mkdir -p "$WORK_DIR/MODULES/$ARCHITECTURE/$PKGNAME"
      cp "$PKG/etc/modulefiles/$PKGNAME" "$WORK_DIR/MODULES/$ARCHITECTURE/$PKGNAME/$PKGVER"
    done < <(find $WORK_DIR/$ARCHITECTURE -maxdepth 2 -mindepth 2 2> /dev/null)
  else
    printf "${EY}WARNING: not updating modulefiles${EZ}\n" >&2
  fi
}

function normModules() {
  echo "$@" | sed -e 's/,/ /g; s/VO_ALICE@//g; s!::!/!g'
}

function existModules() {
  local MODULE
  for MODULE in "$@"; do
    $MODULECMD bash -t avail 2>&1 | grep -qFx "$MODULE" || \
      { printf "${ER}ERROR: $MODULE was not found${EZ}\n" >&2; return 1; }
  done
}

function stripDyld() {
  echo unset $( (typeset | grep ^DYLD | cut -d= -f1 | grep '_modshare$' | xargs echo ) 2> /dev/null; echo ';' )
}

ARCHITECTURE=
WORK_DIR=
NO_REFRESH=
DEFAULT_WORK_DIRS=( "$ALIBUILD_WORK_DIR" "$ALICE_WORK_DIR" "$ALIBUILD_CHDIR/sw" sw ../sw )
ARGS=()
COMMAND_IN_ENV=()
CLEAN_ENV=1
VERBOSE=1

while [[ $# -gt 0 ]]; do
  case "$1" in
    --architecture|-a) ARCHITECTURE="$2"; shift 2 ;;
    --work-dir|-w) WORK_DIR="$2"; shift 2 ;;
    --no-refresh) NO_REFRESH=1; shift ;;
    -q) unset VERBOSE; shift ;;
    --shellrc) unset CLEAN_ENV; shift ;;
    --help|help) printHelp; exit 0 ;;
    -c) shift; COMMAND_IN_ENV=("$@"); break ;;
    *) ARGS+=("$1"); shift ;;
  esac
done

ACTION="${ARGS[0]}"
ARGS=("${ARGS[@]:1}")
export ALIENVLVL=$((ALIENVLVL+1))

if [[ "$ACTION" == "shell-helper" ]]; then
cat <<\EOF
# Source this file from your .bashrc, .zshrc or .kshrc.
ALIBUILD_ALIENV__="$ALIBUILD_ALIENV"
if test -z "$ALIBUILD_ALIENV__"; then
  ALIBUILD_ALIENV__=`unset alienv; export LC_ALL=C; export LANG=C; type -p alienv | sed -e 's/^alienv is //g' 2> /dev/null`
fi
alienv() {
  if test -z "$ALIBUILD_ALIENV__"; then
    echo "alienv not found"
    return 1
  fi
  for A in "$@"; do
    if ! test `echo "$A" | cut -b1` = -; then
      if test "$A" = load || test "$A" = unload; then
        unset A
        eval `"$ALIBUILD_ALIENV__" "$@"`
        return $?
      else
        break
      fi
    elif test "$A" = -c; then
      break
    fi
  done
  unset A
  "$ALIBUILD_ALIENV__" "$@"
}
EOF
  exit 0
fi

if [[ -z "$WORK_DIR" ]]; then
  for WORK_DIR in "${DEFAULT_WORK_DIRS[@]}"; do
    [[ -d "$WORK_DIR" ]] && break || WORK_DIR=
  done
  [[ -z "$WORK_DIR" ]] && { printHelp "No default work dir can be accessed:"                      \
                                      "export \$ALIBUILD_WORK_DIR or run alienv from a directory" \
                                      "containing the \"sw\" dir"; false; }
fi
[[ ! -d "$WORK_DIR" ]] && { printHelp "Work dir $WORK_DIR cannot be accessed"; false; }
WORK_DIR=$(cd "$WORK_DIR"; pwd)
[[ -z "$ARCHITECTURE" ]] && ARCHITECTURE="$("$(dirname "$0")/aliBuild" architecture 2> /dev/null || true)"
[[ -z "$ARCHITECTURE" || "$ARCHITECTURE" == "<unknown>" ]] && ARCHITECTURE=$(ls -1t $WORK_DIR | grep -vE '^[A-Z]+$' | head -n1)
[[ -z "$ARCHITECTURE" ]] && { printHelp "Cannot autodetect architecture"; false; }

# Look for modulecmd (v3) or modulecmd-compat (>= v4)
MODULECMD=$(which modulecmd 2> /dev/null || true)
[[ -x "$MODULECMD" ]] || MODULECMD="$(dirname $(which envml 2> /dev/null || true) 2> /dev/null || true)/../libexec/modulecmd-compat"
[[ -x "$MODULECMD" ]] || MODULECMD="$(brew --prefix modules 2> /dev/null || true)/libexec/modulecmd-compat"
[[ -x "$MODULECMD" ]] || { installHint; false; }

if [[ -d $WORK_DIR/MODULES/$ARCHITECTURE ]]; then
  touch $WORK_DIR/MODULES/$ARCHITECTURE/.testwrite 2> /dev/null || NO_REFRESH=1
  rm -f $WORK_DIR/MODULES/$ARCHITECTURE/.testwrite 2> /dev/null || true
fi

if type csrutil &> /dev/null; then
  csrutil status | grep -q disabled || \
    { printf "${EY}WARNING: System Integrity Protection is enabled. " >&2;
      printf "You might run into environment problems if you do not disable it.${EZ}\n" >&2; }
fi

# Ask environment-modules to run in 3.x mode
export MODULES_USE_COMPAT_VERSION=1

export MODULEPATH="$WORK_DIR/MODULES/$ARCHITECTURE${MODULEPATH:+":$MODULEPATH"}"
MODULEPATH=$(echo "$MODULEPATH" | sed -e 's/::*/:/g; s/^://; s/:$//')
IGNORE_ERR="Unable to locate a modulefile for 'Toolchain/"
case "$ACTION" in
  enter)
    [[ $ALIENVLVL == 1 ]] || \
      { printf "${ER}ERROR: already in an alienv environment${EZ}\n" >&2; exit 1; }
    MODULES=$(normModules "${ARGS[@]}")
    collectModules
    existModules $MODULES
    eval $($MODULECMD bash add $MODULES 2> >(grep -v "$IGNORE_ERR" >&2))
    [[ $UNAME == Darwin ]] && eval $(stripDyld)
    detectShell
    if [[ ! -z "$CLEAN_ENV" ]]; then
      case $MODULES_SHELL in
        sh|bash) export PS1="[$MODULES]"' \w $> '   ;;
        ksh)     export PS1="[$MODULES]"' $PWD $> ' ;;
        zsh)     export PS1="[$MODULES]"' %~ %#> '  ;;
      esac
    fi
    $MODULECMD bash list
    printf "${ET}Use ${EM}alienv list${ET} to list loaded modules. " >&2
    printf "Use ${EM}exit${ET} to exit this environment.${EZ}\n" >&2
    exec ${CLEAN_ENV:+env $SHELL_NORC_ENV} $MODULES_SHELL ${CLEAN_ENV:+$SHELL_NORC_PARAM} -i
  ;;
  printenv|load|unload)
    [[ $ACTION == printenv ]] && ACTION=load
    MODULES=$(normModules "${ARGS[@]}")
    [[ $ACTION == load ]] && { collectModules; existModules $MODULES; }
    [[ $VERBOSE == 1 ]] && printf "${ET}Use ${EM}alienv list${ET} to list loaded modules.${EZ}\n" >&2
    detectShell
    $MODULECMD $MODULES_SHELL $ACTION $MODULES 2> >(grep -v "$IGNORE_ERR" >&2)
    [[ $UNAME == Darwin ]] && ( eval $($MODULECMD $MODULES_SHELL $ACTION $MODULES 2> /dev/null) &> /dev/null; stripDyld )
    exit 0
  ;;
  setenv)
    [[ -z "${COMMAND_IN_ENV[*]}" ]] && { printHelp "No command specified with -c"; false; }
    MODULES=$(normModules "${ARGS[@]}")
    collectModules
    existModules $MODULES
    eval $($MODULECMD bash add $MODULES 2> >(grep -v "$IGNORE_ERR" >&2))
    [[ $UNAME == Darwin ]] && eval $(stripDyld)
    exec "${COMMAND_IN_ENV[@]}"
  ;;
  q|query)
    [[ -z "${ARGS[0]}" ]] && SEARCH_CMD=cat || SEARCH_CMD=( grep -iE "${ARGS[0]}" )
    collectModules
    $MODULECMD bash -t avail 2>&1 | grep -E '^[^/]+/[^/]+$' | grep -vE ':$' | \
      "${SEARCH_CMD[@]}" | sed -e 's!^\([^/]*\)/\(.*\)$!VO_ALICE@\1::\2!'
    exit ${PIPESTATUS[3]}  # grep
  ;;
  avail)
    collectModules
    exec $MODULECMD bash avail
  ;;
  list)
    exec $MODULECMD bash list
  ;;
  modulecmd)
    collectModules
    exec $MODULECMD "${ARGS[@]}"
  ;;
  '')
    printHelp "What do you want to do?"
  ;;
  *)
    printHelp "Unknown command: $1"
  ;;
esac
false