From 76ec36cd244cfb7f15e4a197888e6a27e5a8c40f Mon Sep 17 00:00:00 2001
From: Michael Krasnyk <michael.krasnyk@gmail.com>
Date: Wed, 19 Jun 2024 23:21:41 +0200
Subject: [PATCH] feat: populate symlink tree for local directories

---
 python/BUILD                     |  7 +++++++
 python/poetry_deps.py            |  5 +++--
 python/py_venv.py                | 30 +++-------------------------
 python/tests/poetry_deps_test.py |  4 +++-
 python/utils.py                  | 34 ++++++++++++++++++++++++++++++++
 5 files changed, 50 insertions(+), 30 deletions(-)
 create mode 100644 python/utils.py

diff --git a/python/BUILD b/python/BUILD
index cc40168..1883b56 100644
--- a/python/BUILD
+++ b/python/BUILD
@@ -3,12 +3,19 @@ py_binary(
     srcs = ["poetry_deps.py"],
     visibility = ["__subpackages__"],
     deps = [
+        ":utils",
         "@rules_poetry_pip//:pkg",
     ],
 )
 
+py_library(
+    name = "utils",
+    srcs = ["utils.py"],
+)
+
 py_binary(
     name = "py_venv",
     srcs = ["py_venv.py"],
     visibility = ["__subpackages__"],
+    deps = [":utils"],
 )
diff --git a/python/poetry_deps.py b/python/poetry_deps.py
index 6cc5276..9efb80c 100644
--- a/python/poetry_deps.py
+++ b/python/poetry_deps.py
@@ -15,6 +15,8 @@
     from pip._vendor.packaging.utils import parse_wheel_filename
     from pip._vendor.packaging.version import InvalidVersion
 
+from python.utils import populate_symlink_tree
+
 _SHA256_PREFIX = "sha256:"
 
 
@@ -44,8 +46,7 @@ def install(args):
         local_package_path = Path(args.source_url[0])
         if local_package_path.is_absolute() and local_package_path.is_dir():
             # Add symbolic links to the local directory
-            for item in local_package_path.iterdir():
-                (output_path / item.name).symlink_to(item)
+            populate_symlink_tree(local_package_path, output_path / local_package_path.name)
             return 0
 
         # Otherwise it is a list of files or URLs
diff --git a/python/py_venv.py b/python/py_venv.py
index 85b5ff1..9ab587f 100644
--- a/python/py_venv.py
+++ b/python/py_venv.py
@@ -1,9 +1,8 @@
 import argparse
-import filecmp
-import os
-import warnings
 from pathlib import Path
 
+from python.utils import populate_symlink_tree
+
 SKIP_SET = {Path("requirements.txt")}
 
 
@@ -16,30 +15,7 @@ def main(argv=None):
     args = parser.parse_args(argv)
 
     for python_path in args.path:
-        if not python_path.exists() or not python_path.is_dir():
-            raise RuntimeError(f"Required Python package directory {python_path} does not exist")
-
-        for directory_path, _, file_names in os.walk(python_path):
-            in_package_directory = Path(os.path.relpath(directory_path, python_path))
-            target_directory = args.target / in_package_directory
-            target_directory.mkdir(parents=True, exist_ok=True)
-            relative_directory = Path(os.path.relpath(directory_path, target_directory))
-
-            for file_name in file_names:
-                if in_package_directory / file_name in SKIP_SET:
-                    continue
-
-                symlink_path = target_directory / file_name
-                target_path = relative_directory / file_name
-                if symlink_path.exists():
-                    if not filecmp.cmp(symlink_path, Path(directory_path) / file_name, shallow=False):
-                        warnings.warn(
-                            f"{symlink_path} already exists and points to {os.path.realpath(symlink_path)}\n"
-                            + f"Skip {target_path} which seems to have different contents"
-                        )
-                    continue
-
-                symlink_path.symlink_to(target_path)
+        populate_symlink_tree(python_path, args.target, SKIP_SET)
 
 
 if __name__ == "__main__":
diff --git a/python/tests/poetry_deps_test.py b/python/tests/poetry_deps_test.py
index f36407d..19195d2 100644
--- a/python/tests/poetry_deps_test.py
+++ b/python/tests/poetry_deps_test.py
@@ -56,8 +56,10 @@ def test_install_from_directory(self):
                 retcode = main.install(args)
                 self.assertEqual(retcode, 0)
 
-                wheels = glob.glob(f"{args.output}/six*")
+                wheels = list(args.output.rglob("six*"))
                 self.assertGreater(len(wheels), 0)
+                self.assertTrue(wheels[0].is_symlink())
+                self.assertEqual(wheels[0].parent.name, Path(tmp_input).name)
 
     def test_install_from_url(self):
         args = InstallArgs()
diff --git a/python/utils.py b/python/utils.py
new file mode 100644
index 0000000..dc1e0e0
--- /dev/null
+++ b/python/utils.py
@@ -0,0 +1,34 @@
+import filecmp
+import os
+import warnings
+from pathlib import Path
+
+
+def populate_symlink_tree(source, target, skip_set=None):
+    if not source.exists() or not source.is_dir():
+        raise RuntimeError(f"Required Python package directory {source} does not exist")
+
+    for directory_path, _, file_names in os.walk(source):
+        in_package_directory = Path(os.path.relpath(directory_path, source))
+        target_directory = target / in_package_directory
+        target_directory.mkdir(parents=True, exist_ok=True)
+        if source.is_absolute():
+            relative_directory = Path(directory_path)
+        else:
+            relative_directory = Path(os.path.relpath(directory_path, target_directory))
+
+        for file_name in file_names:
+            if skip_set and in_package_directory / file_name in skip_set:
+                continue
+
+            symlink_path = target_directory / file_name
+            target_path = relative_directory / file_name
+            if symlink_path.exists():
+                if not filecmp.cmp(symlink_path, Path(directory_path) / file_name, shallow=False):
+                    warnings.warn(
+                        f"{symlink_path} already exists and points to {os.path.realpath(symlink_path)}\n"
+                        + f"Skip {target_path} which seems to have different contents"
+                    )
+                continue
+
+            symlink_path.symlink_to(target_path)