From 43efa46bb31f714be651442f309367c280e8537c Mon Sep 17 00:00:00 2001
From: Harsh Modi <hmodi@aviatrix.com>
Date: Wed, 30 Aug 2023 17:37:26 +0000
Subject: [PATCH] pth: Add a `add_pth` option to py_pyenv.

Currently, rules_pyvenv installs all the python deps by symlinking them in site-packages. However, in bazel, the python deps are actually added to the PYTHONPATH, which means that for some deps, just symlinking into site-packages doesn't work.

However, venv supports this via pth files. pth files will append the
directories specified in the pth file to the PYTHONPATH, which exhibits
the behavior we want.

Fixes #13
---
 build_env.py | 13 ++++++++++---
 venv.bzl     |  5 ++++-
 2 files changed, 14 insertions(+), 4 deletions(-)

diff --git a/build_env.py b/build_env.py
index 522e90c..b78932a 100644
--- a/build_env.py
+++ b/build_env.py
@@ -166,13 +166,20 @@ def install_site_file(site_packages_path: pathlib.Path, file: EnvFile) -> None:
         site_path.symlink_to(file.path.resolve())
 
 
-def install_files(env_path: pathlib.Path, files: List[EnvFile]) -> None:
+def install_files(env_path: pathlib.Path, files: List[EnvFile], add_pth: bool) -> None:
     site_packages_path = find_site_packages(env_path)
+    pth = site_packages_path / "venv.pth"
+    pths = set()
     for file in files:
         if is_data_file(file):
             install_data_file(env_path, file)
         else:
-            install_site_file(site_packages_path, file)
+            if add_pth:
+                pths.add(file.env_path.parts[0])
+            else:
+                install_site_file(site_packages_path, file)
+    if add_pth:
+        pth.write_text("\n".join(pths), encoding="utf-8")
 
 
 # A copy of importlib_metadata:entry_points that takes a list of search paths.
@@ -266,7 +273,7 @@ def main():
     builder = venv.EnvBuilder(clear=True, symlinks=True, with_pip=True)
     builder.create(str(env_path))
 
-    install_files(env_path, files)
+    install_files(env_path, files, build_env_input.get("add_pth", False))
     generate_console_scripts(env_path)
 
     extra_commands = build_env_input.get("commands")
diff --git a/venv.bzl b/venv.bzl
index 775e3d4..fcd745b 100644
--- a/venv.bzl
+++ b/venv.bzl
@@ -44,6 +44,7 @@ def _py_venv_deps_impl(ctx):
         "workspace": ctx.workspace_name,
         "imports": imports,
         "files": files,
+        "add_pth": ctx.attr.add_pth,
         "commands": ctx.attr.commands,
         "always_link": ctx.attr.always_link,
     }
@@ -58,12 +59,13 @@ _py_venv_deps = rule(
         "data": attr.label_list(),
         "commands": attr.string_list(),
         "always_link": attr.bool(),
+        "add_pth": attr.bool(),
         "output": attr.output(),
     },
     toolchains = [PYTHON_TOOLCHAIN_TYPE],
 )
 
-def py_venv(name, deps = None, data = None, extra_pip_commands = None, always_link = False, venv_location = None, **kwargs):
+def py_venv(name, deps = None, data = None, extra_pip_commands = None, always_link = False, venv_location = None, add_pth = False, **kwargs):
     deps = deps or []
     data = data or []
     extra_pip_commands = extra_pip_commands or []
@@ -78,6 +80,7 @@ def py_venv(name, deps = None, data = None, extra_pip_commands = None, always_li
         commands = extra_pip_commands,
         always_link = always_link,
         output = out_name,
+        add_pth = add_pth,
         **kwargs,
     )