2 files changed, 48 insertions, 0 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
+