From 8a09ab3da30e4b7f544157c5e90562595814f803 Mon Sep 17 00:00:00 2001
From: Simon Mavi Stewart <simon_stewart@apple.com>
Date: Fri, 6 Dec 2024 16:12:08 +0000
Subject: [PATCH] Add libdir to library search path

---
 python/private/pypi/attrs.bzl       |  6 ++++++
 python/private/pypi/whl_library.bzl | 25 +++++++++++++++++++++----
 2 files changed, 27 insertions(+), 4 deletions(-)

diff --git a/python/private/pypi/attrs.bzl b/python/private/pypi/attrs.bzl
index c9b7ea66a9..32cc3fcce3 100644
--- a/python/private/pypi/attrs.bzl
+++ b/python/private/pypi/attrs.bzl
@@ -15,6 +15,12 @@
 "common attributes for whl_library and pip_repository"
 
 ATTRS = {
+    "add_libdir_to_library_search_path": attr.bool(
+        default = False,
+        doc = """
+If true, add the lib dir of the bundled interpreter to the library search path via LDFLAGS.
+""",
+    ),
     "download_only": attr.bool(
         doc = """
 Whether to use "pip download" instead of "pip wheel". Disables building wheels from source, but allows use of
diff --git a/python/private/pypi/whl_library.bzl b/python/private/pypi/whl_library.bzl
index 612ca2cfdf..7c3289e93a 100644
--- a/python/private/pypi/whl_library.bzl
+++ b/python/private/pypi/whl_library.bzl
@@ -139,11 +139,28 @@ def _parse_optional_attrs(rctx, args, extra_pip_args = None):
     if rctx.attr.enable_implicit_namespace_pkgs:
         args.append("--enable_implicit_namespace_pkgs")
 
+    env = {}
     if rctx.attr.environment != None:
-        args += [
-            "--environment",
-            json.encode(struct(arg = rctx.attr.environment)),
-        ]
+        for key, value in rctx.attr.environment.items():
+            env[key] = value
+
+    # This is super hacky, but working out something nice is tricky.
+    # This is in particular needed for psycopg2 which attempts to link libpython.a,
+    # in order to point the linker at the correct python intepreter.
+    if rctx.attr.add_libdir_to_library_search_path:
+        if "LDFLAGS" in env:
+            fail("Can't set both environment LDFLAGS and add_libdir_to_library_search_path")
+        command = [pypi_repo_utils.resolve_python_interpreter(rctx), "-c", "import sys ; sys.stdout.write('{}/lib'.format(sys.exec_prefix))"]
+        result = rctx.execute(command)
+        if result.return_code != 0:
+            fail("Failed to get LDFLAGS path: command: {}, exit code: {}, stdout: {}, stderr: {}".format(command, result.return_code, result.stdout, result.stderr))
+        libdir = result.stdout
+        env["LDFLAGS"] = "-L{}".format(libdir)
+
+    args += [
+        "--environment",
+        json.encode(struct(arg = env)),
+    ]
 
     return args