Unsource ROS 1 & 2 workspaces from environment variables

unsource_ros.bash

# SPDX-License-Identifier: BSD-3-Clause
# SPDX-FileCopyrightText: Czech Technical University in Prague

###########################
# ROS 1 & 2 Bash Unsource #
###########################

# The provided function `unsource_ros` removes all remnants of ROS 2 from your current environment.
# First source this file (that does nothing) and then execute `unsource_ros` .

remove_path_entry() {
    local verbose="$1"
    local var_name="$2"
    local path_to_remove="$3"

    local current_value="${!var_name}"
    if [[ ":$current_value:" != *":$path_to_remove:"* ]]; then
        return
    fi

    # Split current value into an array and delete all occurrences of $current_value
    local -a PATH_ARRAY
    IFS=':' read -ra PATH_ARRAY <<< "$current_value"
    local i
    for i in "${!PATH_ARRAY[@]}"; do
      if [[ "${PATH_ARRAY[i]}" = "$path_to_remove" ]]; then
        [[ $verbose -lt 2 ]] || echo "${var_name}: Remove path ${PATH_ARRAY[i]}"
        unset 'PATH_ARRAY[i]'
      fi
    done

    local new_path="$(IFS=:; echo "${PATH_ARRAY[*]}")"
    export "$var_name"="$new_path"
}

remove_path_prefix() {
    local verbose="$1"
    local var_name="$2"
    local prefix="$3"

    local current_value="${!var_name}"
    if [[ ":$current_value" != *":$prefix"* ]]; then
        return
    fi

    # Split current value into an array and delete all occurrences of $prefix
    local -a PATH_ARRAY
    IFS=':' read -ra PATH_ARRAY <<< "$current_value"
    local i
    for i in "${!PATH_ARRAY[@]}"; do
      if [[ "${PATH_ARRAY[i]}" = "$prefix"* ]]; then
        [[ $verbose -lt 2 ]] || echo "${var_name}: Remove prefix ${PATH_ARRAY[i]}"
        unset 'PATH_ARRAY[i]'
      fi
    done

    local new_path="$(IFS=:; echo "${PATH_ARRAY[*]}")"
    export "$var_name"="$new_path"
}

process_dsv_file() {
    local verbose="$1"
    local dsv_file="$2"
    local install_prefix="$3"

    # There are multiple cases what the DSV prefix might be

    process_dsv_file_in_prefix "$verbose" "$dsv_file" "$install_prefix"
    process_dsv_file_in_prefix "$verbose" "$dsv_file" "${install_prefix}/share"

    # Also try to find subfolder <prefix>/<pkg> belonging to this DSV
    IFS='/' read -ra PREFIX_PARTS <<< "$install_prefix"
    IFS='/' read -ra DSV_PARTS <<< "$dsv_file"
    process_dsv_file_in_prefix "$verbose" "$dsv_file" "$install_prefix/${DSV_PARTS[${#PREFIX_PARTS[@]}]}"
}

process_dsv_file_in_prefix() {
    local verbose="$1"
    local dsv_file="$2"
    local install_prefix="$3"

    local line
    while IFS= read -r line; do
      process_dsv_line_in_prefix "$verbose" "$dsv_file" "$install_prefix" "$line"
    done < "$dsv_file"
}

process_dsv_line_in_prefix() {
    local verbose="$1"
    local dsv_file="$2"
    local install_prefix="$3"
    local line="$4"

    local -a PARTS
    IFS=';' read -ra PARTS <<< "$line"
    local op="${PARTS[0]}"
    local var="${PARTS[1]}"
    local val="${PARTS[2]}"

    # Per docs: "If the value is not an absolute path the prefix path is prepended."
    # Per docs: "An empty value therefore represents the prefix path."
    local full_path="$val"
    if [[ -z "$val" ]]; then
        full_path="$install_prefix"
    elif [[ "$val" != /* ]]; then
        full_path="$install_prefix/$val"
    fi

    case "$op" in
        prepend-non-duplicate|prepend-non-duplicate-if-exists)
            remove_path_entry "$verbose" "$var" "$full_path"
            ;;
        set|set-if-unset)
            [[ $verbose -lt 1 ]] || echo "Unsetting '${var}' !"
            unset "$var"
            ;;
        source)
            # Nothing to do here
            ;;
    esac
}

var_cleanup() {
    local var_name="$1"
    if [ ! -v "$var_name" ]; then
        return
    fi

    local val="${!var_name}"
    # Replace double colons with single
    val="${val//::/:}"
    # Remove leading colon
    val="${val#:}"
    # Remove trailing colon
    val="${val%:}"
    export "$var_name"="$val"
}

function unsource_ros() {
  args=( "$@" )

  local use_dsv=0
  local verbose=0
  local extra=0
  local custom_envs=0
  local help=0

  local i
  for i in "${!args[@]}"; do
    case "${args[i]}" in
        --dsv|-d)
            use_dsv=1
            unset 'args[i]'
            ;;
        --verbose|-v)
            verbose=$((verbose+1))
            unset 'args[i]'
            ;;
        -vv)
            verbose=$((verbose+2))
            unset 'args[i]'
            ;;
        --extra|--extra-env|-e)
            extra=1
            unset 'args[i]'
            ;;
        --help|-h)
            help=1
            unset 'args[i]'
            ;;
      esac
  done

  [[ $verbose -lt 2 ]] || echo "STEP: Arguments parsed"

  if [ "$help" = "1" ]; then
    echo "Usage: unsource_ros [--verbose|v] [--dsv|-d] [-h] [WS_PATH [...]]"
    echo
    echo -e "\tWS_PATH\t\tIf no paths are given, unsource all workspaces. If at least one is given, only unsource the given workspaces."
    echo -e "\t--verbose|-v\tTurn on info prints (you may pass this options twice to get more)"
    echo -e "\t--dsv|-d\tAlso unsource everything specified by DSV env files. This is slower, but more accurate."
    echo -e "\t--extra-env|-e\tAlso unset non-path ROS variables like ROS_DISTRO etc."
    echo -e "\t--help|-h\tShow this help message and exit."
    return
  fi

  local -a PREFIXES

  [[ $verbose -lt 2 ]] || echo "STEP: Detecting workspaces"

  if [[ ${#args[@]} -gt 0 ]]; then
    local arg
    for arg in "${args[@]}"; do
      if [ -d "$arg" ]; then
        # normalize the path, but do not use realpath to retain symlinks
        PREFIXES+=("$(cd "$arg"; pwd)")
      else
        echo "Prefix $arg does not exist, ignoring it!" >&2
      fi
    done
    custom_envs=1
  else
    local -a ROS1_PREFIXES
    IFS=':' read -ra ROS1_PREFIXES <<< "$ROS_PACKAGE_PATH"

    local -a ROS2_PREFIXES
    IFS=':' read -ra ROS2_PREFIXES <<< "$COLCON_PREFIX_PATH"

    local -a CMAKE_PREFIXES
    IFS=':' read -ra CMAKE_PREFIXES <<< "$CMAKE_PREFIX_PATH"

    PREFIXES=("${ROS1_PREFIXES[@]}" "${ROS2_PREFIXES[@]}" "${CMAKE_PREFIXES[@]}")

    local ros_distro
    while IFS= read -r ros_distro; do
      if [[ ":$COLCON_PREFIX_PATH:" != *":$ros_distro:"* ]]; then
        PREFIXES+=("$ros_distro")
      fi
    done < <(find "/opt/ros" -maxdepth 1 -mindepth 1 -type d 2>/dev/null)
  fi

  if [[ ${#PREFIXES[@]} -eq 0 ]]; then
    echo "No workspace paths given, exiting." >&2
    return 1
  else
    if [[ $verbose -gt 1 ]]; then
      echo "Detected workspaces:"
      local prefix
      for prefix in "${PREFIXES[@]}"; do
        echo -e "\t${prefix}"
      done
    fi
  fi

  # Remember LD_LIBRARY_PATH before we ran, to be able to calculate how many paths we have pruned
  local -a LD_LIB_PATH_ORIG
  IFS=':' read -ra LD_LIB_PATH_ORIG <<< "$LD_LIBRARY_PATH"

  [[ $verbose -lt 2 ]] || echo "STEP: Unset whole path variables"

  # First, unset the ROS-only variables where we don't need to keep anything
  local var
  for var in ${!ROS@} ${!AMENT@} ${!COLCON@} ${!CATKIN@}; do
    if [[ "$var" = *"_PATH" ]]; then
      if [[ $custom_envs -eq 0 ]]; then
        [[ $verbose -lt 1 ]] || echo "Unsetting $var"
        unset "$var"
      fi
    elif [[ $extra -eq 1 ]]; then
        [[ $verbose -lt 1 ]] || echo "Unsetting $var"
        unset "$var"
    fi
  done

  [[ $verbose -lt 2 ]] || echo "STEP: Fast unsource"

  # Fast unsource just by prefix matching in a few well-known variables
  local vars=("PATH" "LD_LIBRARY_PATH" "PYTHONPATH" "CMAKE_PREFIX_PATH" "PKG_CONFIG_PATH")
  if [[ $custom_envs -eq 1 ]]; then
    # if unsourcing all workspaces (custom_envs == 0), we have already unset these paths as a whole
    vars+=("ROS_PACKAGE_PATH" "AMENT_PREFIX_PATH" "COLCON_PREFIX_PATH")
  fi
  local prefix
  for prefix in "${PREFIXES[@]}"; do
      [[ $verbose -lt 1 ]] || echo "Unsourcing $prefix"
      local var
      for var in "${vars[@]}"; do
          remove_path_prefix "$verbose" "$var" "$prefix"
      done
  done

  [[ $verbose -lt 2 ]] || echo "STEP: DSV unsource"

  # Slow unsource (crawl all DSV files and undo their effects)
  if [[ $use_dsv -eq 1 ]]; then
    local prefix
    for prefix in "${PREFIXES[@]}"; do
        if [ -d "$prefix" ]; then
            [[ $verbose -lt 1 ]] || echo "Unsourcing $prefix via DSV files"
            local dsv
            while IFS= read -r dsv; do
                process_dsv_file "$verbose" "$dsv" "$prefix"
            done < <(find "$prefix" -path "*/environment/*.dsv")
        fi
    done
  fi

  [[ $verbose -lt 2 ]] || echo "STEP: Cleanup"

  var_cleanup "PATH"
  var_cleanup "LD_LIBRARY_PATH"
  var_cleanup "PYTHONPATH"
  var_cleanup "CMAKE_PREFIX_PATH"
  var_cleanup "ROS_PACKAGE_PATH"
  var_cleanup "COLCON_PREFIX_PATH"
  var_cleanup "AMENT_PREFIX_PATH"

  [[ $verbose -lt 2 ]] || echo "STEP: Done"

  local -a LD_LIB_PATH_NEW
  IFS=':' read -ra LD_LIB_PATH_NEW <<< "$LD_LIBRARY_PATH"

  [[ $verbose -lt 1 ]] || echo "ROS is gone (LD_LIBRARY_PATH from ${#LD_LIB_PATH_ORIG[@]} down to ${#LD_LIB_PATH_NEW[@]} paths)!"
}