From e3395daced2bcecd4b5e385a56604bc67e71b43c Mon Sep 17 00:00:00 2001
From: Danilo Reyes <danilo.reyes.251@proton.me>
Date: Sat, 20 Dec 2025 18:11:36 -0600
Subject: [PATCH] Enhance nix-shell shebang handling in flake.nix by supporting
 multiple interpreters and improving package resolution for Python scripts.
 Introduced logic to create proper Python wrappers and separate system
 packages from Python packages for better environment management.

---
 flake.nix | 73 ++++++++++++++++++++++++++++++++++++++++++++++---------
 1 file changed, 62 insertions(+), 11 deletions(-)

diff --git a/flake.nix b/flake.nix
index 7116eb2..7076034 100644
--- a/flake.nix
+++ b/flake.nix
@@ -58,31 +58,82 @@
       };
       packages.x86_64-linux =
         let
-          # Handle nix-shell shebangs by extracting packages and using writeShellApplication
+          # Handle nix-shell shebangs by extracting packages and interpreter
           scriptBin = path: name:
             let
               content = builtins.readFile path;
-              # Match the nix-shell package line: #! nix-shell -i bash -p package1 package2 ...
-              nixShellMatch = builtins.match ".*#! nix-shell -i [^ ]+ -p ([^\n]+).*" content;
+              # Match the nix-shell line: #! nix-shell -i interpreter -p package1 package2 ...
+              nixShellMatch = builtins.match ".*#! nix-shell -i ([^ ]+) -p ([^\n]+).*" content;
             in
             if nixShellMatch != null then
               let
-                packagesStr = builtins.head nixShellMatch;
+                interpreter = builtins.head nixShellMatch;
+                packagesStr = builtins.elemAt nixShellMatch 1;
                 # Split by spaces and filter empty strings
                 packages = builtins.filter (s: s != "") (pkgs.lib.splitString " " packagesStr);
+                # Resolve package references - handle python3Packages.* and regular packages
+                resolvePackage = pkgName:
+                  if pkgs.lib.hasPrefix "python3Packages." pkgName then
+                    let
+                      pythonPkgName = pkgs.lib.removePrefix "python3Packages." pkgName;
+                    in
+                    pkgs.python3Packages.${pythonPkgName}
+                  else if pkgName == "python3" then
+                    pkgs.python3
+                  else
+                    pkgs.${pkgName};
                 # Get package references from pkgs
-                runtimeInputs = builtins.map (pkgName: pkgs.${pkgName}) packages;
+                runtimeInputs = builtins.map resolvePackage packages;
                 # Remove the nix-shell shebang lines
                 scriptContent = builtins.replaceStrings
-                  [ "#!/usr/bin/env nix-shell\n" "#! nix-shell -i bash -p ${packagesStr}\n" ]
+                  [ "#!/usr/bin/env nix-shell\n" "#! nix-shell -i ${interpreter} -p ${packagesStr}\n" ]
                   [ "" "" ]
                   content;
               in
-              pkgs.writeShellApplication {
-                inherit name;
-                runtimeInputs = runtimeInputs;
-                text = scriptContent;
-              }
+              if interpreter == "python3" then
+                # For Python scripts, create a proper wrapper with Python shebang
+                let
+                  # Separate Python packages from system packages
+                  pythonPkgNames = builtins.filter (pkg: pkgs.lib.hasPrefix "python3Packages." pkg) packages;
+                  systemPkgNames = builtins.filter (pkg: !(pkgs.lib.hasPrefix "python3Packages." pkg) && pkg != "python3") packages;
+                  # Resolve Python packages for python3.withPackages
+                  pythonPkgs = builtins.map (pkgName:
+                    let pythonPkgName = pkgs.lib.removePrefix "python3Packages." pkgName;
+                    in pkgs.python3Packages.${pythonPkgName}
+                  ) pythonPkgNames;
+                  # Create Python environment with required packages
+                  pythonEnv = if pythonPkgs != [] then
+                    pkgs.python3.withPackages (ps: pythonPkgs)
+                  else
+                    pkgs.python3;
+                  # Resolve system packages for PATH
+                  systemRuntimeInputs = builtins.map resolvePackage systemPkgNames;
+                  scriptFile = pkgs.writeText "${name}.py" ''
+                    #!${pythonEnv}/bin/python3
+                    ${scriptContent}
+                  '';
+                in
+                pkgs.stdenv.mkDerivation {
+                  inherit name;
+                  buildInputs = [ pkgs.makeWrapper ];
+                  dontUnpack = true;
+                  installPhase = ''
+                    mkdir -p $out/bin
+                    cp ${scriptFile} $out/bin/${name}
+                    chmod +x $out/bin/${name}
+                    ${pkgs.lib.optionalString (systemRuntimeInputs != []) ''
+                      wrapProgram $out/bin/${name} \
+                        --prefix PATH : ${pkgs.lib.makeBinPath systemRuntimeInputs}
+                    ''}
+                  '';
+                }
+              else
+                # For shell scripts, use writeShellApplication
+                pkgs.writeShellApplication {
+                  inherit name;
+                  runtimeInputs = runtimeInputs;
+                  text = scriptContent;
+                }
             else
               pkgs.writeScriptBin name content;
           pkgsBin =