From 0ce4c4ceea7ddd24dc5452aa0df5be63cf6c5907 Mon Sep 17 00:00:00 2001
From: Raz <raz.rotenberg@gmail.com>
Date: Wed, 30 Nov 2022 17:47:56 +0200
Subject: [PATCH] Supporting zsh

* Parameter substitution is done differently for zsh
* Using @ modifier to split strings

- Added environment variable 'GENV_SHELL' (based on $SHELL)
- Added helper function '_genv_get_env'
---
 docs/source/development/development.rst |  2 +-
 docs/source/development/reference.rst   |  7 +++++++
 libexec/genv                            |  1 +
 libexec/genv-sh-activate                | 26 ++++++++++++++++++-------
 libexec/genv-sh-deactivate              | 15 ++++++++++++--
 5 files changed, 41 insertions(+), 10 deletions(-)

diff --git a/docs/source/development/development.rst b/docs/source/development/development.rst
index 1c3052c..b74e023 100644
--- a/docs/source/development/development.rst
+++ b/docs/source/development/development.rst
@@ -19,7 +19,7 @@ Open a terminal in the project directory.
 
 .. warning::
 
-   We support :code:`bash` terminals only at the moment. Query your shell using :code:`echo $SHELL` and run :code:`bash` if it is not your shell, or use a :ref:`Docker <Docker>` setup instead.
+   We support :code:`bash` and :code:`zsh` terminals only at the moment. Query your shell using :code:`echo $SHELL` and run :code:`bash` if your shell is not supported, or use a :ref:`Docker <Docker>` setup instead.
 
 Then run the following commands to configure your terminal:
 
diff --git a/docs/source/development/reference.rst b/docs/source/development/reference.rst
index ed1c581..ae4f38e 100644
--- a/docs/source/development/reference.rst
+++ b/docs/source/development/reference.rst
@@ -30,6 +30,13 @@ This is mainly for development environments on CPU machines where :code:`nvidia-
 
 ----
 
+:code:`GENV_SHELL`
+
+Set the shell type (e.g. `zsh`, `bash`).
+Default is based on environment variable :code:`$SHELL`.
+
+----
+
 :code:`GENV_TMPDIR`
 
 Path of the temp directory where all state JSON files are stored.
diff --git a/libexec/genv b/libexec/genv
index cc639b1..2ff2d71 100755
--- a/libexec/genv
+++ b/libexec/genv
@@ -4,6 +4,7 @@ set -e
 [ -n "$GENV_DEBUG" ] && set -x
 
 export GENV_ROOT=${GENV_ROOT:-"$(cd -- "$(dirname -- "${BASH_SOURCE[0]}")" &> /dev/null && pwd)/.."}
+export GENV_SHELL=${GENV_SHELL:-${SHELL##*/}} # TODO(raz): we should check the actual running shell and not the default one
 export PATH="$GENV_ROOT/libexec:$PATH"
 export PYTHONPATH="$GENV_ROOT/py:$PYTHONPATH"
 
diff --git a/libexec/genv-sh-activate b/libexec/genv-sh-activate
index 0b84fa7..81bb35d 100755
--- a/libexec/genv-sh-activate
+++ b/libexec/genv-sh-activate
@@ -126,10 +126,20 @@ done
 function functions()
 {
   cat <<EOS
+_genv_get_env()
+{
+  if [ "$GENV_SHELL" = "zsh" ]; then
+    echo "\${(P)1}"
+  else
+    echo "\${!1}"
+  fi
+}
+
 _genv_append_to_env()
 {
   # based on https://unix.stackexchange.com/a/415028
-  export \$1="\${!1:+\${!1}:}\$2"
+  val=\$(_genv_get_env \$1)
+  export \$1="\${val:+\$val:}\$2"
 }
 
 _genv_set_env()
@@ -146,19 +156,21 @@ _genv_unset_env()
 
 _genv_backup_env()
 {
-  if [ -n "\${!1}" ]; then
-    export GENV_BACKUP_ENV_\$1="\${!1}"
+  val=\$(_genv_get_env \$1)
+  if [ -n "\$val" ]; then
+    export GENV_BACKUP_ENV_\$1="\$val"
     _genv_append_to_env GENV_BACKUP_ENVS \$1
   fi
 }
 
 _genv_restore_env()
 {
-  backup="GENV_BACKUP_ENV_\$1"
-  if [ -n "\${!backup}" ]; then
-    export \$1="\${!backup}"
+  var="GENV_BACKUP_ENV_\$1"
+  val=\$(_genv_get_env \$var)
+  if [ -n "\$val" ]; then
+    export \$1="\$val"
   fi
-  unset \$backup
+  unset \$var
 }
 
 _genv_replace_env()
diff --git a/libexec/genv-sh-deactivate b/libexec/genv-sh-deactivate
index 56e7493..164da88 100755
--- a/libexec/genv-sh-deactivate
+++ b/libexec/genv-sh-deactivate
@@ -30,7 +30,12 @@ function functions()
   cat <<EOS
 _genv_unset_envs()
 {
-  IFS=: read -a names <<< "\$GENV_ENVS"
+  if [ "$GENV_SHELL" = "zsh" ]; then
+    names=("\${(@s/:/)GENV_ENVS}")
+  else
+    IFS=: read -a names <<< "\$GENV_ENVS"
+  fi
+
   unset GENV_ENVS
 
   for name in "\${names[@]}"
@@ -41,7 +46,12 @@ _genv_unset_envs()
 
 _genv_restore_envs()
 {
-  IFS=: read -a names <<< "\$GENV_BACKUP_ENVS"
+  if [ "$GENV_SHELL" = "zsh" ]; then
+    names=("\${(@s/:/)GENV_BACKUP_ENVS}")
+  else
+    IFS=: read -a names <<< "\$GENV_BACKUP_ENVS"
+  fi
+
   unset GENV_BACKUP_ENVS
 
   for name in "\${names[@]}"
@@ -71,6 +81,7 @@ unset -f _genv_backup_env
 unset -f _genv_unset_env
 unset -f _genv_set_env
 unset -f _genv_append_to_env
+unset -f _genv_get_env
 EOS
 }