From 17785242cb334d4003052d1ed14a6bb00840cdb4 Mon Sep 17 00:00:00 2001
From: Chan Lee <Leetimemp@gmail.com>
Date: Tue, 20 May 2025 10:59:09 +0800
Subject: [PATCH] stubgen: add `LIB_PATH` option to locate dependent shared
 libraries

---
 cmake/nanobind-config.cmake |  6 +++++-
 src/stubgen.py              | 21 +++++++++++++++++++++
 2 files changed, 26 insertions(+), 1 deletion(-)

diff --git a/cmake/nanobind-config.cmake b/cmake/nanobind-config.cmake
index 8f8a3ae9b..1eb172a82 100644
--- a/cmake/nanobind-config.cmake
+++ b/cmake/nanobind-config.cmake
@@ -562,7 +562,7 @@ endfunction()
 # ---------------------------------------------------------------------------
 
 function (nanobind_add_stub name)
-  cmake_parse_arguments(PARSE_ARGV 1 ARG "VERBOSE;INCLUDE_PRIVATE;EXCLUDE_DOCSTRINGS;INSTALL_TIME;EXCLUDE_FROM_ALL" "MODULE;OUTPUT;MARKER_FILE;COMPONENT;PATTERN_FILE" "PYTHON_PATH;DEPENDS")
+  cmake_parse_arguments(PARSE_ARGV 1 ARG "VERBOSE;INCLUDE_PRIVATE;EXCLUDE_DOCSTRINGS;INSTALL_TIME;EXCLUDE_FROM_ALL" "MODULE;OUTPUT;MARKER_FILE;COMPONENT;PATTERN_FILE" "PYTHON_PATH;LIB_PATH;DEPENDS")
 
   if (EXISTS ${NB_DIR}/src/stubgen.py)
     set(NB_STUBGEN "${NB_DIR}/src/stubgen.py")
@@ -590,6 +590,10 @@ function (nanobind_add_stub name)
     list(APPEND NB_STUBGEN_ARGS -i "${TMP}")
   endforeach()
 
+  foreach (TMP IN LISTS ARG_LIB_PATH)
+    list(APPEND NB_STUBGEN_ARGS -L "${TMP}")
+  endforeach()
+
   if (ARG_PATTERN_FILE)
     list(APPEND NB_STUBGEN_ARGS -p "${ARG_PATTERN_FILE}")
   endif()
diff --git a/src/stubgen.py b/src/stubgen.py
index 5bb9bf3b8..5436f8a9f 100755
--- a/src/stubgen.py
+++ b/src/stubgen.py
@@ -1210,6 +1210,16 @@ def parse_options(args: List[str]) -> argparse.Namespace:
         help="add the directory to the Python import path (can specify multiple times)",
     )
 
+    parser.add_argument(
+        "-L",
+        "--lib-path",
+        action="append",
+        metavar="PATH",
+        dest="lib_paths",
+        default=[],
+        help="add directory to shared library search path (can specify multiple times)"
+    )
+
     parser.add_argument(
         "-m",
         "--module",
@@ -1342,6 +1352,7 @@ def add_pattern(query: str, lines: List[str]):
 
 def main(args: Optional[List[str]] = None) -> None:
     import sys
+    import os
 
     # Ensure that the current directory is on the path
     if "" not in sys.path and "." not in sys.path:
@@ -1362,6 +1373,16 @@ def main(args: Optional[List[str]] = None) -> None:
     for i in opt.imports:
         sys.path.insert(0, i)
 
+    if os.name == 'nt':
+        for lib_path in opt.lib_paths:
+            os.add_dll_directory(lib_path)
+    else:
+        lib_env = "DYLD_LIBRARY_PATH" if sys.platform == "darwin" else "LD_LIBRARY_PATH"
+        old_value = os.environ.get(lib_env, "")
+        paths_str = ":".join(opt.lib_paths)
+        if paths_str:
+            os.environ[lib_env] = f"{paths_str}:{old_value}" if old_value else paths_str
+
     for i, mod in enumerate(opt.modules):
         if not opt.quiet:
             if i > 0: