From 8db990e559d20ea0e261b153c4658c8ae7d8ba48 Mon Sep 17 00:00:00 2001
From: Giel van Schijndel <giel@mortis.eu>
Date: Fri, 12 Apr 2024 14:11:29 +0200
Subject: [PATCH 1/2] zsh: add test for programs.zsh.localVariables

---
 tests/modules/programs/zsh/default.nix        |  1 +
 .../modules/programs/zsh/local-variables.nix  | 30 +++++++++++++++++++
 2 files changed, 31 insertions(+)
 create mode 100644 tests/modules/programs/zsh/local-variables.nix

diff --git a/tests/modules/programs/zsh/default.nix b/tests/modules/programs/zsh/default.nix
index 25aa3b470..479238649 100644
--- a/tests/modules/programs/zsh/default.nix
+++ b/tests/modules/programs/zsh/default.nix
@@ -1,4 +1,5 @@
 {
+  zsh-local-variables = ./local-variables.nix;
   zsh-session-variables = ./session-variables.nix;
   zsh-history-path-new-default = ./history-path-new-default.nix;
   zsh-history-path-new-custom = ./history-path-new-custom.nix;
diff --git a/tests/modules/programs/zsh/local-variables.nix b/tests/modules/programs/zsh/local-variables.nix
new file mode 100644
index 000000000..e22ad1894
--- /dev/null
+++ b/tests/modules/programs/zsh/local-variables.nix
@@ -0,0 +1,30 @@
+{ lib, ... }:
+
+with lib;
+
+{
+  config = {
+    programs.zsh = {
+      enable = true;
+
+      localVariables = rec {
+        V1 = true;
+        V2 = false;
+        V3 = "some-string";
+        V4 = 42;
+        V5 = [ V1 V2 V3 V4 ];
+      };
+    };
+
+    test.stubs.zsh = { };
+
+    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")'
+      '';
+  };
+}

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 2/2] 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 c6901350f..771a7c274 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 c4520e7a7..c152249a3 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 e22ad1894..111cc5fe0 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")''
+      }
+    '';
   };
 }