From c67b06ff923a4e7eb5c93fe0277a966c894afdc4 Mon Sep 17 00:00:00 2001
From: Giel van Schijndel <giel@mortis.eu>
Date: Fri, 12 Apr 2024 14:14:03 +0200
Subject: [PATCH] zsh: support attrsets as associative arrays in
 programs.zsh.localVariables

---
 modules/lib/zsh.nix                           | 11 ++++++++-
 modules/programs/zsh.nix                      |  5 +++-
 .../modules/programs/zsh/local-variables.nix  | 24 +++++++++++++------
 3 files changed, 31 insertions(+), 9 deletions(-)

diff --git a/modules/lib/zsh.nix b/modules/lib/zsh.nix
index c6901350..771a7c27 100644
--- a/modules/lib/zsh.nix
+++ b/modules/lib/zsh.nix
@@ -9,11 +9,20 @@ rec {
       ''"${v}"''
     else if builtins.isList v then
       "(${lib.concatStringsSep " " (map toZshValue v)})"
+    else if builtins.isAttrs v then
+      "(${
+        lib.concatStringsSep " "
+        (lib.mapAttrsToList (n: v: "[${lib.escapeShellArg n}]=${toZshValue v}")
+          v)
+      })"
     else
       ''"${toString v}"'';
 
   # Produces a Zsh shell like definition statement
-  define = n: v: "${n}=${toZshValue v}";
+  define = n: v:
+    "${lib.optionalString (builtins.isAttrs v) "typeset -A "}${n}=${
+      toZshValue v
+    }";
 
   # Given an attribute set containing shell variable names and their
   # assignments, this function produces a string containing a definition
diff --git a/modules/programs/zsh.nix b/modules/programs/zsh.nix
index c4520e7a..c152249a 100644
--- a/modules/programs/zsh.nix
+++ b/modules/programs/zsh.nix
@@ -501,7 +501,10 @@ in
       localVariables = mkOption {
         type = types.attrs;
         default = {};
-        example = { POWERLEVEL9K_LEFT_PROMPT_ELEMENTS=["dir" "vcs"]; };
+        example = {
+          POWERLEVEL9K_LEFT_PROMPT_ELEMENTS=["dir" "vcs"];
+          ZSH_HIGHLIGHT_STYLES={unknown-token="bg=red,fg=white,bold";};
+        };
         description = ''
           Extra local variables defined at the top of {file}`.zshrc`.
         '';
diff --git a/tests/modules/programs/zsh/local-variables.nix b/tests/modules/programs/zsh/local-variables.nix
index e22ad189..111cc5fe 100644
--- a/tests/modules/programs/zsh/local-variables.nix
+++ b/tests/modules/programs/zsh/local-variables.nix
@@ -12,7 +12,13 @@ with lib;
         V2 = false;
         V3 = "some-string";
         V4 = 42;
-        V5 = [ V1 V2 V3 V4 ];
+        V5 = builtins.attrValues V6;
+        V6 = {
+          a = V1;
+          b = V2;
+          c = V3;
+          d = V4;
+        };
       };
     };
 
@@ -20,11 +26,15 @@ with lib;
 
     nmt.script = ''
       assertFileExists home-files/.zshrc
-      assertFileContains home-files/.zshrc 'V1=true'
-      assertFileContains home-files/.zshrc 'V2=false'
-      assertFileContains home-files/.zshrc 'V3="some-string"'
-      assertFileContains home-files/.zshrc 'V4="42"'
-      assertFileContains home-files/.zshrc 'V5=(true false "some-string" "42")'
-      '';
+      assertFileRegex home-files/.zshrc '^V1=true$'
+      assertFileRegex home-files/.zshrc '^V2=false$'
+      assertFileRegex home-files/.zshrc '^V3="some-string"$'
+      assertFileRegex home-files/.zshrc '^V4="42"$'
+      assertFileRegex home-files/.zshrc '^V5=[(]true false "some-string" "42"[)]$'
+      assertFileContains home-files/.zshrc ${
+        lib.escapeShellArg ''
+          typeset -A V6=(['a']=true ['b']=false ['c']="some-string" ['d']="42")''
+      }
+    '';
   };
 }