From db96098311ed09f1fdbe8c624854ab01ea9b525d Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?L=C3=A1szl=C3=B3=20Vask=C3=B3?=
 <1771332+vlaci@users.noreply.github.com>
Date: Tue, 29 Oct 2024 17:24:39 +0100
Subject: [PATCH] python: remove `LD_LIBRARY_PATH` hack from running python
 environment

Injecting `LD_LIBRARY_PATH` to the Python runtime environment is great
to bypass the need of having to patch non-nix binaries loaded into
that environment, however it breaks down, when Python executes any
other program not compiled for the given Nix system, e.g. a shell
script via `subprocess`.

To work this around, `devenv` will inject a `pth`[^1] file to the virtual
environment it creates, which mangles the `LD_LIBRARY_PATH` variable,
undoing any changes to it made by `devenv` but preserving changes from
other sources.

Fixes #1111

[^1]: https://docs.python.org/3/library/site.html
---
 src/modules/languages/python.nix              | 27 +++++++++++++++++++
 tests/python-library-path-poetry/.test.sh     | 11 ++++++++
 tests/python-library-path-poetry/devenv.nix   |  9 +++++++
 .../python-library-path-poetry/pyproject.toml |  9 +++++++
 tests/python-library-path-uv/.test.sh         | 13 +++++++++
 tests/python-library-path-uv/devenv.nix       | 10 +++++++
 tests/python-library-path/.test.sh            | 11 ++++++++
 tests/python-library-path/devenv.nix          |  9 +++++++
 8 files changed, 99 insertions(+)
 create mode 100755 tests/python-library-path-poetry/.test.sh
 create mode 100644 tests/python-library-path-poetry/devenv.nix
 create mode 100644 tests/python-library-path-poetry/pyproject.toml
 create mode 100755 tests/python-library-path-uv/.test.sh
 create mode 100644 tests/python-library-path-uv/devenv.nix
 create mode 100755 tests/python-library-path/.test.sh
 create mode 100644 tests/python-library-path/devenv.nix

diff --git a/src/modules/languages/python.nix b/src/modules/languages/python.nix
index c1b7195be..03411b93c 100644
--- a/src/modules/languages/python.nix
+++ b/src/modules/languages/python.nix
@@ -14,6 +14,9 @@ let
     python = cfg.package;
     requiredPythonModules = cfg.package.pkgs.requiredPythonModules;
     makeWrapperArgs = [
+      "--set"
+      "DEVENV_LD_LIBRARY_PATH_PREFIX"
+      libraries
       "--prefix"
       "LD_LIBRARY_PATH"
       ":"
@@ -39,6 +42,20 @@ let
     follows = [ "nixpkgs" ];
   };
 
+  pth_file =
+    let
+      exec_content = ''
+        ld_library_path = os.environ.get("LD_LIBRARY_PATH")
+        ld_library_path_prefix = os.environ.get("DEVENV_LD_LIBRARY_PATH_PREFIX")
+        if ld_library_path and ld_library_path_prefix:
+            if ld_library_path == ld_library_path_prefix:
+                del os.environ["LD_LIBRARY_PATH"]
+            else:
+                os.environ["LD_LIBRARY_PATH"] = ld_library_path.removeprefix(ld_library_path_prefix + ":")
+      '';
+    in
+    pkgs.writeText "devenv.pth" ''import os; exec("""${builtins.replaceStrings [ "\n" ] [ "\\n" ] exec_content}""")'';
+
   initVenvScript =
     let
       USE_UV_SYNC = cfg.uv.sync.enable && builtins.compareVersions cfg.uv.package.version "0.4.4" >= 0;
@@ -80,6 +97,9 @@ let
           ''
         }
         echo "${package.interpreter}" > "$VENV_PATH/.devenv_interpreter"
+        ${lib.optionalString pkgs.stdenv.isLinux ''
+          ln -snf ${pth_file} "$VENV_PATH/${package.sitePackages}/devenv.pth"
+        ''}
       fi
 
       source "$VENV_PATH"/bin/activate
@@ -157,6 +177,9 @@ let
         if "''${UV_SYNC_COMMAND[@]}"
         then
           echo "$ACTUAL_UV_CHECKSUM" > "$UV_CHECKSUM_FILE"
+          ${lib.optionalString pkgs.stdenv.isLinux ''
+            ln -snf ${pth_file} "$VENV_PATH/${package.sitePackages}/devenv.pth"
+          ''}
         else
           echo "uv sync failed. Run 'uv sync' manually." >&2
           exit 1
@@ -186,6 +209,10 @@ let
 
       # Make sure poetry's venv uses the configured Python executable.
       ${cfg.poetry.package}/bin/poetry env use --no-interaction --quiet ${package.interpreter}
+
+      ${lib.optionalString pkgs.stdenv.isLinux ''
+        ln -snf ${pth_file} ".venv/${package.sitePackages}/devenv.pth"
+      ''}
     }
 
     function _devenv_poetry_install
diff --git a/tests/python-library-path-poetry/.test.sh b/tests/python-library-path-poetry/.test.sh
new file mode 100755
index 000000000..dbd0a6e5b
--- /dev/null
+++ b/tests/python-library-path-poetry/.test.sh
@@ -0,0 +1,11 @@
+python - <<EOF
+import os
+
+assert "LD_LIBRARY_PATH" not in os.environ
+EOF
+
+LD_LIBRARY_PATH=set-from-shell python - <<EOF
+import os
+
+assert os.environ["LD_LIBRARY_PATH"] == "set-from-shell"
+EOF
diff --git a/tests/python-library-path-poetry/devenv.nix b/tests/python-library-path-poetry/devenv.nix
new file mode 100644
index 000000000..4bc16ff09
--- /dev/null
+++ b/tests/python-library-path-poetry/devenv.nix
@@ -0,0 +1,9 @@
+{ pkgs, ... }:
+
+{
+  languages.python = {
+    enable = true;
+    poetry.enable = true;
+    libraries = [ pkgs.file ];
+  };
+}
diff --git a/tests/python-library-path-poetry/pyproject.toml b/tests/python-library-path-poetry/pyproject.toml
new file mode 100644
index 000000000..3bec1af6c
--- /dev/null
+++ b/tests/python-library-path-poetry/pyproject.toml
@@ -0,0 +1,9 @@
+[tool.poetry]
+name = "python-ld-library-path"
+version = "0.1.0"
+description = ""
+authors = ["Your Name <you@example.com>"]
+
+[build-system]
+requires = ["poetry-core"]
+build-backend = "poetry.core.masonry.api"
diff --git a/tests/python-library-path-uv/.test.sh b/tests/python-library-path-uv/.test.sh
new file mode 100755
index 000000000..1976f1245
--- /dev/null
+++ b/tests/python-library-path-uv/.test.sh
@@ -0,0 +1,13 @@
+python - <<EOF
+import os
+
+assert "LD_LIBRARY_PATH" not in os.environ
+EOF
+
+
+LD_LIBRARY_PATH=set-from-shell python - <<EOF
+import os
+import sys
+print(sys.prefix)
+assert os.environ["LD_LIBRARY_PATH"] == "set-from-shell"
+EOF
diff --git a/tests/python-library-path-uv/devenv.nix b/tests/python-library-path-uv/devenv.nix
new file mode 100644
index 000000000..e88c747c0
--- /dev/null
+++ b/tests/python-library-path-uv/devenv.nix
@@ -0,0 +1,10 @@
+{ pkgs, ... }:
+
+{
+  languages.python = {
+    enable = true;
+    venv.enable = true;
+    uv.enable = true;
+    libraries = [ pkgs.file ];
+  };
+}
diff --git a/tests/python-library-path/.test.sh b/tests/python-library-path/.test.sh
new file mode 100755
index 000000000..dbd0a6e5b
--- /dev/null
+++ b/tests/python-library-path/.test.sh
@@ -0,0 +1,11 @@
+python - <<EOF
+import os
+
+assert "LD_LIBRARY_PATH" not in os.environ
+EOF
+
+LD_LIBRARY_PATH=set-from-shell python - <<EOF
+import os
+
+assert os.environ["LD_LIBRARY_PATH"] == "set-from-shell"
+EOF
diff --git a/tests/python-library-path/devenv.nix b/tests/python-library-path/devenv.nix
new file mode 100644
index 000000000..d8fc5c5d2
--- /dev/null
+++ b/tests/python-library-path/devenv.nix
@@ -0,0 +1,9 @@
+{ pkgs, ... }:
+
+{
+  languages.python = {
+    enable = true;
+    venv.enable = true;
+    libraries = [ pkgs.file ];
+  };
+}