From b04aa56503c59d07a72361dfe9845c3ab2faac1b Mon Sep 17 00:00:00 2001
From: Martin Toman <toman.martin@live.com>
Date: Sun, 1 Nov 2020 18:27:47 +0100
Subject: [PATCH] powerline-go: add zsh integration

---
 modules/programs/powerline-go.nix             | 35 +++++++++++++++----
 .../powerline-go/{standard.nix => bash.nix}   |  0
 .../modules/programs/powerline-go/default.nix |  5 ++-
 tests/modules/programs/powerline-go/zsh.nix   | 31 ++++++++++++++++
 4 files changed, 63 insertions(+), 8 deletions(-)
 rename tests/modules/programs/powerline-go/{standard.nix => bash.nix} (100%)
 create mode 100644 tests/modules/programs/powerline-go/zsh.nix

diff --git a/modules/programs/powerline-go.nix b/modules/programs/powerline-go.nix
index a4cd233cf..8f5db8f26 100644
--- a/modules/programs/powerline-go.nix
+++ b/modules/programs/powerline-go.nix
@@ -106,17 +106,38 @@ in {
     };
   };
 
-  config = mkIf (cfg.enable && config.programs.bash.enable) {
-    programs.bash.initExtra = ''
-      function _update_ps1() {
-        local old_exit_status=$?
-        PS1="$(${pkgs.powerline-go}/bin/powerline-go -error $old_exit_status ${commandLineArguments})"
+  config = {
+    programs.bash.initExtra =
+      mkIf (cfg.enable && config.programs.bash.enable) ''
+        function _update_ps1() {
+          local old_exit_status=$?
+          PS1="$(${pkgs.powerline-go}/bin/powerline-go -error $old_exit_status ${commandLineArguments})"
+          ${cfg.extraUpdatePS1}
+          return $old_exit_status
+        }
+
+        if [ "$TERM" != "linux" ]; then
+          PROMPT_COMMAND="_update_ps1;$PROMPT_COMMAND"
+        fi
+      '';
+
+    programs.zsh.initExtra = mkIf (cfg.enable && config.programs.zsh.enable) ''
+      function powerline_precmd() {
+        PS1="$(${pkgs.powerline-go}/bin/powerline-go -error $? -shell zsh ${commandLineArguments})"
         ${cfg.extraUpdatePS1}
-        return $old_exit_status
+      }
+
+      function install_powerline_precmd() {
+        for s in "$\{precmd_functions[@]}"; do
+          if [ "$s" = "powerline_precmd" ]; then
+            return
+          fi
+        done
+        precmd_functions+=(powerline_precmd)
       }
 
       if [ "$TERM" != "linux" ]; then
-        PROMPT_COMMAND="_update_ps1;$PROMPT_COMMAND"
+        install_powerline_precmd
       fi
     '';
   };
diff --git a/tests/modules/programs/powerline-go/standard.nix b/tests/modules/programs/powerline-go/bash.nix
similarity index 100%
rename from tests/modules/programs/powerline-go/standard.nix
rename to tests/modules/programs/powerline-go/bash.nix
diff --git a/tests/modules/programs/powerline-go/default.nix b/tests/modules/programs/powerline-go/default.nix
index 50febefbc..e89d7c6cd 100644
--- a/tests/modules/programs/powerline-go/default.nix
+++ b/tests/modules/programs/powerline-go/default.nix
@@ -1 +1,4 @@
-{ powerline-go-standard = ./standard.nix; }
+{
+  powerline-go-bash = ./bash.nix;
+  powerline-go-zsh = ./zsh.nix;
+}
diff --git a/tests/modules/programs/powerline-go/zsh.nix b/tests/modules/programs/powerline-go/zsh.nix
new file mode 100644
index 000000000..702149528
--- /dev/null
+++ b/tests/modules/programs/powerline-go/zsh.nix
@@ -0,0 +1,31 @@
+{ config, lib, pkgs, ... }:
+
+with lib;
+
+{
+  config = {
+    programs = {
+      zsh.enable = true;
+
+      powerline-go = {
+        enable = true;
+        newline = true;
+        modules = [ "nix-shell" ];
+        pathAliases = { "\\~/project/foo" = "prj-foo"; };
+        settings = {
+          ignore-repos = [ "/home/me/project1" "/home/me/project2" ];
+        };
+      };
+    };
+
+    nixpkgs.overlays =
+      [ (self: super: { zsh = pkgs.writeScriptBin "dummy-zsh" ""; }) ];
+
+    nmt.script = ''
+      assertFileExists home-files/.zshrc
+      assertFileContains \
+        home-files/.zshrc \
+        '/bin/powerline-go -error $? -shell zsh -modules nix-shell -newline -path-aliases \~/project/foo=prj-foo -ignore-repos /home/me/project1,/home/me/project2'
+    '';
+  };
+}