From 2c8e5d157f9d1751e88df2ebdc2ff6730cae1a12 Mon Sep 17 00:00:00 2001
From: "ax.vivien" <ax.vivien@criteo.com>
Date: Mon, 6 Nov 2023 17:33:57 +0100
Subject: [PATCH] also add packages that are placed in src folder

---
 cluster_pack/packaging.py                     |  5 ++++-
 tests/test_packaging.py                       | 22 +++++++++++++++++--
 tests/user-lib-src-layout/__init__.py         |  0
 tests/user-lib-src-layout/setup.py            |  8 +++++++
 .../src/user_lib/__init__.py                  |  0
 .../user-lib-src-layout/src/user_lib/test.py  |  2 ++
 .../src/user_lib2/__init__.py                 |  0
 7 files changed, 34 insertions(+), 3 deletions(-)
 create mode 100644 tests/user-lib-src-layout/__init__.py
 create mode 100644 tests/user-lib-src-layout/setup.py
 create mode 100644 tests/user-lib-src-layout/src/user_lib/__init__.py
 create mode 100644 tests/user-lib-src-layout/src/user_lib/test.py
 create mode 100644 tests/user-lib-src-layout/src/user_lib2/__init__.py

diff --git a/cluster_pack/packaging.py b/cluster_pack/packaging.py
index ec68cfb..92bdcf3 100644
--- a/cluster_pack/packaging.py
+++ b/cluster_pack/packaging.py
@@ -327,7 +327,10 @@ def _get_editable_requirements(executable: str = sys.executable) -> List[str]:
     top_level_pkgs = []
     for pkg in _get_packages(True, executable):
         location = pkg.get("editable_project_location", pkg.get("location", ""))
-        for _pkg in setuptools.find_packages(location):
+        packages_found = set(
+            setuptools.find_packages(location) + setuptools.find_packages(f"{location}/src")
+            )
+        for _pkg in packages_found:
             if "." in _pkg:
                 continue
             imported = __import__(_pkg)
diff --git a/tests/test_packaging.py b/tests/test_packaging.py
index c70eb24..329694a 100644
--- a/tests/test_packaging.py
+++ b/tests/test_packaging.py
@@ -64,6 +64,19 @@ def test__get_editable_requirements():
         assert "user_lib2" in pkg_names
 
 
+def test__get_editable_requirements_for_src_layout():
+    with tempfile.TemporaryDirectory() as tempdir:
+        _create_venv(tempdir)
+        _pip_install(tempdir)
+        editable_requirements = packaging._get_editable_requirements(
+            f"{tempdir}/bin/python", use_src_layout=True
+            )
+        assert len(editable_requirements) == 2
+        pkg_names = [os.path.basename(req) for req in editable_requirements]
+        assert "user_lib" in pkg_names
+        assert "user_lib2" in pkg_names
+
+
 @pytest.mark.skipif(sys.version_info < (3, 7), reason="requires python3.7 or higher")
 def test__get_editable_requirements_withpip23():
     with tempfile.TemporaryDirectory() as tempdir:
@@ -90,10 +103,11 @@ def _create_venv(tempdir: str):
     subprocess.check_call([sys.executable, "-m", "venv", f"{tempdir}"])
 
 
-def _pip_install(tempdir: str, pip_version: str = "18.1"):
+def _pip_install(tempdir: str, pip_version: str = "18.1", use_src_layout: bool = False):
     subprocess.check_call([f"{tempdir}/bin/python", "-m", "pip", "install",
                            "cloudpickle", f"pip=={pip_version}"])
-    pkg = _get_editable_package_name()
+    pkg = (_get_editable_package_name_src_layout() if use_src_layout
+           else _get_editable_package_name())
     subprocess.check_call([f"{tempdir}/bin/python", "-m", "pip", "install", "-e", pkg])
     if pkg not in sys.path:
         sys.path.append(pkg)
@@ -103,6 +117,10 @@ def _get_editable_package_name():
     return os.path.join(os.path.dirname(__file__), "user-lib")
 
 
+def _get_editable_package_name_src_layout():
+    return os.path.join(os.path.dirname(__file__), "user-lib-src-layout")
+
+
 def test_get_current_pex_filepath():
     with tempfile.TemporaryDirectory() as tempdir:
         path_to_pex = f"{tempdir}/out.pex"
diff --git a/tests/user-lib-src-layout/__init__.py b/tests/user-lib-src-layout/__init__.py
new file mode 100644
index 0000000..e69de29
diff --git a/tests/user-lib-src-layout/setup.py b/tests/user-lib-src-layout/setup.py
new file mode 100644
index 0000000..1e720e5
--- /dev/null
+++ b/tests/user-lib-src-layout/setup.py
@@ -0,0 +1,8 @@
+#!/usr/bin/env python3
+
+import setuptools
+
+setuptools.setup(
+    name='user_lib_src_layout',
+    version='0.0.1',
+    packages=setuptools.find_packages())
diff --git a/tests/user-lib-src-layout/src/user_lib/__init__.py b/tests/user-lib-src-layout/src/user_lib/__init__.py
new file mode 100644
index 0000000..e69de29
diff --git a/tests/user-lib-src-layout/src/user_lib/test.py b/tests/user-lib-src-layout/src/user_lib/test.py
new file mode 100644
index 0000000..0c91099
--- /dev/null
+++ b/tests/user-lib-src-layout/src/user_lib/test.py
@@ -0,0 +1,2 @@
+def test_method():
+    print("Hello World!")
diff --git a/tests/user-lib-src-layout/src/user_lib2/__init__.py b/tests/user-lib-src-layout/src/user_lib2/__init__.py
new file mode 100644
index 0000000..e69de29