python: add pythonNamespacesHook

4 files changed, 68 insertions, 1 deletions

diff --git a/pkgs/development/interpreters/python/hooks/default.nix b/pkgs/development/interpreters/python/hooks/default.nix
index 47690320e81..4d736426f3b 100644
--- a/pkgs/development/interpreters/python/hooks/default.nix
+++ b/pkgs/development/interpreters/python/hooks/default.nix
@@ -89,6 +89,14 @@ in rec {
       };
     } ./python-imports-check-hook.sh) {};
 
+  pythonNamespacesHook = callPackage ({}:
+    makeSetupHook {
+      name = "python-namespaces-hook.sh";
+      substitutions = {
+        inherit pythonSitePackages;
+      };
+    } ./python-namespaces-hook.sh) {};
+
   pythonRemoveBinBytecodeHook = callPackage ({ }:
     makeSetupHook {
       name = "python-remove-bin-bytecode-hook";
diff --git a/pkgs/development/interpreters/python/hooks/python-namespaces-hook.sh b/pkgs/development/interpreters/python/hooks/python-namespaces-hook.sh
new file mode 100644
index 00000000000..50f21819d17
--- /dev/null
+++ b/pkgs/development/interpreters/python/hooks/python-namespaces-hook.sh
@@ -0,0 +1,40 @@
+# Clean up __init__.py's found in namespace directories
+echo "Sourcing python-namespaces-hook"
+
+pythonNamespacesHook() {
+    echo "Executing pythonNamespacesHook"
+
+    for namespace in ${pythonNamespaces[@]}; do
+        echo "Enforcing PEP420 namespace: ${namespace}"
+
+        # split namespace into segments. "azure.mgmt" -> "azure mgmt"
+        IFS='.' read -ra pathSegments <<< $namespace
+        constructedPath=$out/@pythonSitePackages@
+
+        # Need to remove the __init__.py at each namespace level
+        # E.g `azure/__init__.py` and `azure/mgmt/__init__.py`
+        # The __pycache__ entry also needs to be removed
+        for pathSegment in ${pathSegments[@]}; do
+            constructedPath=${constructedPath}/${pathSegment}
+            pathToRemove=${constructedPath}/__init__.py
+            pycachePath=${constructedPath}/__pycache__/__init__*
+
+            if [ -f "$pathToRemove" ]; then
+                echo "Removing $pathToRemove"
+                rm "$pathToRemove"
+            fi
+
+            if [ -f "$pycachePath" ]; then
+                echo "Removing $pycachePath"
+                rm "$pycachePath"
+            fi
+        done
+    done
+
+    echo "Finished executing pythonNamespacesHook"
+}
+
+if [ -z "${dontUsePythonNamespacesHook-}" -a -n "${pythonNamespaces-}" ]; then
+    postFixupHooks+=(pythonNamespacesHook)
+fi
+
diff --git a/pkgs/development/interpreters/python/mk-python-derivation.nix b/pkgs/development/interpreters/python/mk-python-derivation.nix
index 180bc63857c..22938a45585 100644
--- a/pkgs/development/interpreters/python/mk-python-derivation.nix
+++ b/pkgs/development/interpreters/python/mk-python-derivation.nix
@@ -16,6 +16,7 @@
 , pipInstallHook
 , pythonCatchConflictsHook
 , pythonImportsCheckHook
+, pythonNamespacesHook
 , pythonRemoveBinBytecodeHook
 , pythonRemoveTestsDirHook
 , setuptoolsBuildHook
@@ -131,6 +132,9 @@ let
   ] ++ lib.optionals (stdenv.buildPlatform == stdenv.hostPlatform) [
     # This is a test, however, it should be ran independent of the checkPhase and checkInputs
     pythonImportsCheckHook
+  ] ++ lib.optionals (python.pythonAtLeast "3.3") [
+    # Optionally enforce PEP420 for python3
+    pythonNamespacesHook
   ] ++ nativeBuildInputs;
 
   buildInputs = buildInputs ++ pythonPath;
diff --git a/pkgs/top-level/python-packages.nix b/pkgs/top-level/python-packages.nix
index 7f9f28b4539..c4551e783c4 100644
--- a/pkgs/top-level/python-packages.nix
+++ b/pkgs/top-level/python-packages.nix
@@ -108,7 +108,22 @@ in {
   inherit buildSetupcfg;
 
   inherit (callPackage ../development/interpreters/python/hooks { })
-    eggUnpackHook eggBuildHook eggInstallHook flitBuildHook pipBuildHook pipInstallHook pytestCheckHook pythonCatchConflictsHook pythonImportsCheckHook pythonRemoveBinBytecodeHook pythonRemoveTestsDirHook setuptoolsBuildHook setuptoolsCheckHook venvShellHook wheelUnpackHook;
+    eggUnpackHook
+    eggBuildHook
+    eggInstallHook
+    flitBuildHook
+    pipBuildHook
+    pipInstallHook
+    pytestCheckHook
+    pythonCatchConflictsHook
+    pythonImportsCheckHook
+    pythonNamespacesHook
+    pythonRemoveBinBytecodeHook
+    pythonRemoveTestsDirHook
+    setuptoolsBuildHook
+    setuptoolsCheckHook
+    venvShellHook
+    wheelUnpackHook;
 
   # helpers