stubgen: add LIB_PATH option to locate dependent shared libraries

Author: Doekin

cmake/nanobind-config.cmake

@@ -590,7 +590,7 @@ endfunction()
 # ---------------------------------------------------------------------------
 
 function (nanobind_add_stub name)
-  cmake_parse_arguments(PARSE_ARGV 1 ARG "VERBOSE;INCLUDE_PRIVATE;EXCLUDE_DOCSTRINGS;EXCLUDE_VALUES;INSTALL_TIME;RECURSIVE;EXCLUDE_FROM_ALL" "MODULE;COMPONENT;PATTERN_FILE;OUTPUT_PATH" "PYTHON_PATH;DEPENDS;MARKER_FILE;OUTPUT")
+  cmake_parse_arguments(PARSE_ARGV 1 ARG "VERBOSE;INCLUDE_PRIVATE;EXCLUDE_DOCSTRINGS;EXCLUDE_VALUES;INSTALL_TIME;RECURSIVE;EXCLUDE_FROM_ALL" "MODULE;COMPONENT;PATTERN_FILE;OUTPUT_PATH" "PYTHON_PATH;LIB_PATH;DEPENDS;MARKER_FILE;OUTPUT")
 
   if (EXISTS ${NB_DIR}/src/stubgen.py)
     set(NB_STUBGEN "${NB_DIR}/src/stubgen.py")
@@ -626,6 +626,10 @@ function (nanobind_add_stub name)
     list(APPEND NB_STUBGEN_ARGS -i "${PYTHON_PATH}")
   endforeach()
 
+  foreach (LIB_PATH IN LISTS ARG_LIB_PATH)
+    list(APPEND NB_STUBGEN_ARGS -L "${LIB_PATH}")
+  endforeach()
+
   if (ARG_PATTERN_FILE)
     list(APPEND NB_STUBGEN_ARGS -p "${ARG_PATTERN_FILE}")
   endif()

src/stubgen.py

@@ -1285,6 +1285,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",
@@ -1426,6 +1436,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:
@@ -1446,6 +1457,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: