From 46409b5c32da64d06ec465d7cb0b2b6a1e1ad246 Mon Sep 17 00:00:00 2001
From: Frederik Rietdijk <fridh@fridh.nl>
Date: Mon, 22 May 2017 11:01:00 +0200
Subject: [PATCH] Python: add sitecustomize.py, listen to NIX_PYTHONPATH

This commit adds a Nix-specific module that recursively adds paths that
are on `NIX_PYTHONPATH` to `sys.path`. In order to process possible
`.pth` files `site.addsitedir` is used.

The paths listed in `PYTHONPATH` are added to `sys.path` afterwards, but
they will be added before the entries we add here and thus take
precedence.

The reason for adding support for this environment variable is that we
can set it in a wrapper without breaking support for `PYTHONPATH`.
---
 .../python/cpython/2.7/default.nix             |  5 +++++
 .../interpreters/python/cpython/default.nix    |  3 +++
 .../interpreters/python/pypy/default.nix       |  3 +++
 .../interpreters/python/pypy/prebuilt.nix      |  4 ++++
 .../interpreters/python/sitecustomize.py       | 18 ++++++++++++++++++
 5 files changed, 33 insertions(+)
 create mode 100644 pkgs/development/interpreters/python/sitecustomize.py

diff --git a/pkgs/development/interpreters/python/cpython/2.7/default.nix b/pkgs/development/interpreters/python/cpython/2.7/default.nix
index 1503c221bdd..de980f1ca68 100644
--- a/pkgs/development/interpreters/python/cpython/2.7/default.nix
+++ b/pkgs/development/interpreters/python/cpython/2.7/default.nix
@@ -256,6 +256,11 @@ in with passthru; stdenv.mkDerivation ({
 
     inherit passthru;
 
+    postFixup = ''
+      # Include a sitecustomize.py file. Note it causes an error when it's in postInstall with 2.7.
+      cp ${../../sitecustomize.py} $out/${sitePackages}/sitecustomize.py
+    '';
+
     enableParallelBuilding = true;
 
     doCheck = false; # expensive, and fails
diff --git a/pkgs/development/interpreters/python/cpython/default.nix b/pkgs/development/interpreters/python/cpython/default.nix
index 5eb7f3ec166..2f398d8dccb 100644
--- a/pkgs/development/interpreters/python/cpython/default.nix
+++ b/pkgs/development/interpreters/python/cpython/default.nix
@@ -221,6 +221,9 @@ in with passthru; stdenv.mkDerivation {
     find $out/lib/python*/config-* -type f -print -exec nuke-refs -e $out '{}' +
     find $out/lib -name '_sysconfigdata*.py*' -print -exec nuke-refs -e $out '{}' +
 
+    # Include a sitecustomize.py file
+    cp ${../sitecustomize.py} $out/${sitePackages}/sitecustomize.py
+
     # Determinism: rebuild all bytecode
     # We exclude lib2to3 because that's Python 2 code which fails
     # We rebuild three times, once for each optimization level
diff --git a/pkgs/development/interpreters/python/pypy/default.nix b/pkgs/development/interpreters/python/pypy/default.nix
index 42b652978bc..705ddb2a6f5 100644
--- a/pkgs/development/interpreters/python/pypy/default.nix
+++ b/pkgs/development/interpreters/python/pypy/default.nix
@@ -137,6 +137,9 @@ in with passthru; stdenv.mkDerivation rec {
 
     # Python on Nix is not manylinux1 compatible. https://github.com/NixOS/nixpkgs/issues/18484
     echo "manylinux1_compatible=False" >> $out/lib/${libPrefix}/_manylinux.py
+
+    # Include a sitecustomize.py file
+    cp ${../sitecustomize.py} $out/${sitePackages}/sitecustomize.py
   '';
 
   inherit passthru;
diff --git a/pkgs/development/interpreters/python/pypy/prebuilt.nix b/pkgs/development/interpreters/python/pypy/prebuilt.nix
index 6af198e65ea..1ac779ed25e 100644
--- a/pkgs/development/interpreters/python/pypy/prebuilt.nix
+++ b/pkgs/development/interpreters/python/pypy/prebuilt.nix
@@ -84,6 +84,10 @@ in with passthru; stdenv.mkDerivation {
     echo "Removing bytecode"
     find . -name "__pycache__" -type d -depth -exec rm -rf {} \;
     popd
+
+    # Include a sitecustomize.py file
+    cp ${../sitecustomize.py} $out/${sitePackages}/sitecustomize.py
+
   '';
 
   doInstallCheck = true;
diff --git a/pkgs/development/interpreters/python/sitecustomize.py b/pkgs/development/interpreters/python/sitecustomize.py
new file mode 100644
index 00000000000..7e43be56147
--- /dev/null
+++ b/pkgs/development/interpreters/python/sitecustomize.py
@@ -0,0 +1,18 @@
+"""
+This is a Nix-specific module for discovering modules built with Nix.
+
+The module recursively adds paths that are on `NIX_PYTHONPATH` to `sys.path`. In
+order to process possible `.pth` files `site.addsitedir` is used.
+
+The paths listed in `PYTHONPATH` are added to `sys.path` afterwards, but they
+will be added before the entries we add here and thus take precedence.
+
+Note the `NIX_PYTHONPATH` environment variable is unset in order to prevent leakage.
+"""
+import site
+import os
+import functools
+
+paths = os.environ.pop('NIX_PYTHONPATH', None)
+if paths:
+    functools.reduce(lambda k, p: site.addsitedir(p, k), paths.split(':'), site._init_pathinfo())