From 9061c3f6aaa50ffdad657c5fcd740d15b533718e Mon Sep 17 00:00:00 2001
From: DavHau <d.hauer.it@gmail.com>
Date: Sat, 31 Aug 2024 17:59:54 +0200
Subject: [PATCH] devshells: improve select-shell

---
 devShell.nix                 |  2 +-
 pkgs/scripts/select-shell.py | 90 ++++++++++++++++++++++++++----------
 2 files changed, 66 insertions(+), 26 deletions(-)

diff --git a/devShell.nix b/devShell.nix
index d615ff71c..c23c49627 100644
--- a/devShell.nix
+++ b/devShell.nix
@@ -38,7 +38,7 @@
         ];
         shellHook = ''
           echo -e "${ansiEscapes.green}switch to another dev-shell using: select-shell${ansiEscapes.reset}"
-          export PROJECT_ROOT=$(git rev-parse --show-toplevel)
+          export PRJ_ROOT=$(git rev-parse --show-toplevel)
         '';
       };
     };
diff --git a/pkgs/scripts/select-shell.py b/pkgs/scripts/select-shell.py
index c2542d9aa..cc8484f3e 100644
--- a/pkgs/scripts/select-shell.py
+++ b/pkgs/scripts/select-shell.py
@@ -1,39 +1,79 @@
 import argparse
 import json
+import os
 import pathlib
 import subprocess
 import sys
 
-parser = argparse.ArgumentParser(description="Select a devshell")
-parser.add_argument("shell", help="the name of the devshell to select", nargs="?")
-parser.add_argument("--list", action="store_true", help="list available devshells")
-args = parser.parse_args()
-
 selected_shell_file = pathlib.Path(".direnv/selected-shell")
 
-if not args.list and not args.shell:
-    parser.print_help()
-    exit(0)
-if args.list:
+
+def parse_args() -> argparse.Namespace:
+    parser = argparse.ArgumentParser(description="Select a devshell")
+    parser.add_argument("shell", help="the name of the devshell to select", nargs="?")
+    parser.add_argument("--list", action="store_true", help="list available devshells")
+    parser.add_argument(
+        "--show", action="store_true", help="show the currently selected devshell"
+    )
+    return parser.parse_args()
+
+
+def get_current_shell() -> str | None:
+    if selected_shell_file.exists():
+        with open(selected_shell_file) as f:
+            return f.read().strip()
+    return None
+
+
+def print_current_shell() -> None:
+    current_shell = get_current_shell()
+    if current_shell:
+        print(f"Currently selected devshell: {current_shell}")
+    else:
+        print("No devshell selected")
+
+
+def print_devshells() -> None:
+    project_root = os.environ.get("PRJ_ROOT")
     flake_show = subprocess.run(
-        ["nix", "flake", "show", "--json", "--no-write-lock-file"],
+        [
+            "nix",
+            "eval",
+            "--json",
+            "--no-write-lock-file",
+            "--apply",
+            "shells: builtins.mapAttrs (name: _shell: name) shells",
+            f"{project_root}#devShells.x86_64-linux",
+        ],
         stdout=subprocess.PIPE,
     )
-    data = json.loads(flake_show.stdout.decode())
-    print("Available devshells:")
-    print("\n".join(data["devShells"]["x86_64-linux"].keys()))
-    exit(0)
-if selected_shell_file.exists():
-    with open(selected_shell_file) as f:
-        current_shell = f.read().strip()
-else:
-    current_shell = ""
+    names = json.loads(flake_show.stdout.decode())
+    print("Available devshells:\n")
+    print("\n".join(names))
 
-if current_shell == args.shell:
-    print(f"{args.shell} devshell already selected. No changes made.")
-    sys.exit(0)
 
-with open(selected_shell_file, "w") as f:
-    f.write(args.shell)
+def select_shell(shell: str) -> None:
+    if shell == get_current_shell():
+        print(f"{shell} devshell already selected. No changes made.")
+        sys.exit(0)
+    else:
+        with open(selected_shell_file, "w") as f:
+            f.write(shell)
+        print(f"{shell} devshell selected")
 
-print(f"{args.shell} devshell selected")
+
+def main() -> None:
+    args = parse_args()
+    if args.show:
+        print_current_shell()
+    elif args.list:
+        print_devshells()
+    elif args.shell:
+        select_shell(args.shell)
+    else:
+        print_current_shell()
+        print("Use --help for more information")
+
+
+if __name__ == "__main__":
+    main()