Merge pull request #419 from Simple-Robotics/update-nanobind

Use nanobind's new recursive stub generation, fix dynamic module handling

Author: jorisv

CHANGELOG.md

@@ -6,11 +6,15 @@ The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.0.0/).
 
 ## [Unreleased]
 
+### Added
+- Recursive stub generation for Python bindings ([#419](https://github.com/Simple-Robotics/proxsuite/pull/419))
+
 ### Changed
 - Change the default branch to `devel` ([#395](https://github.com/Simple-Robotics/proxsuite/pull/395))
 - Change `dual_feasibility` test threshold in `sparse_maros_meszaros` unit test ([#403](https://github.com/Simple-Robotics/proxsuite/pull/403))
 - replace `std::numeric_limits<T>::infinity()` by `std::numeric_limits<T>::max()` ([#413](https://github.com/Simple-Robotics/proxsuite/pull/413))
 - Upgraded nanobind dependency version (submodule) to v2.9.2 ([#418](https://github.com/Simple-Robotics/proxsuite/pull/418))
+- Better dynamic module handling ([#419](https://github.com/Simple-Robotics/proxsuite/pull/419))
 
 ### Fixed
 - Use the right table to store `configure-args` cmeel argument ([#403](https://github.com/Simple-Robotics/proxsuite/pull/403))

bindings/python/CMakeLists.txt

@@ -39,6 +39,7 @@ if(NOT nanobind_FOUND)
     external/nanobind
     ${CMAKE_CURRENT_BINARY_DIR}/external/nanobind
   )
+  message(STATUS "Will use nanobind submodule.")
 else()
   message(STATUS "Found installed nanobind.")
 endif()
@@ -107,12 +108,18 @@ function(list_filter list regular_expression dest_list)
   set(${dest_list} ${list} PARENT_SCOPE)
 endfunction(list_filter)
 
-function(CREATE_PYTHON_TARGET target_name COMPILE_OPTIONS dependencies)
+function(
+  create_python_target
+  target_name
+  compile_options
+  dependencies
+  generate_stubs
+)
   nanobind_add_module(${target_name} ${PYWRAP_SOURCES} ${PYWRAP_HEADERS})
   add_dependencies(${PROJECT_NAME}_python ${target_name})
 
   target_link_libraries(${target_name} PUBLIC ${dependencies})
-  target_compile_options(${target_name} PRIVATE ${COMPILE_OPTIONS})
+  target_compile_options(${target_name} PRIVATE ${compile_options})
   target_link_libraries(${target_name} PRIVATE proxsuite)
   target_compile_definitions(
     ${target_name}
@@ -133,7 +140,7 @@ function(CREATE_PYTHON_TARGET target_name COMPILE_OPTIONS dependencies)
     if(LINK_PYTHON_INTERFACE_TO_OPENMP)
       target_link_libraries(${target_name} PRIVATE ${OpenMP_CXX_LIBRARIES})
     endif(LINK_PYTHON_INTERFACE_TO_OPENMP)
-  else(BUILD_WITH_OPENMP_SUPPORT)
+  else()
     list_filter("${PYWRAP_HEADERS}" "expose-parallel" PYWRAP_HEADERS)
   endif(BUILD_WITH_OPENMP_SUPPORT)
 
@@ -173,16 +180,26 @@ function(CREATE_PYTHON_TARGET target_name COMPILE_OPTIONS dependencies)
   endif()
 
   install(TARGETS ${target_name} DESTINATION ${${PYWRAP}_INSTALL_DIR})
-  if(GENERATE_PYTHON_STUBS)
+  if(${generate_stubs})
+    set(
+      stub_outputs
+      ${target_name}/__init__.pyi
+      ${target_name}/helpers.pyi
+      ${target_name}/proxqp/__init__.pyi
+      ${target_name}/proxqp/dense.pyi
+      ${target_name}/proxqp/sparse.pyi
+    )
     nanobind_add_stub(
       ${target_name}_stub
       MODULE ${target_name}
-      OUTPUT ${target_name}.pyi
+      OUTPUT_PATH ${CMAKE_CURRENT_BINARY_DIR}
+      OUTPUT ${stub_outputs}
+      RECURSIVE
       PYTHON_PATH $<TARGET_FILE_DIR:${target_name}>
       DEPENDS ${target_name}
     )
     install(
-      FILES ${CMAKE_CURRENT_BINARY_DIR}/${target_name}.pyi
+      DIRECTORY ${CMAKE_CURRENT_BINARY_DIR}/${target_name}
       DESTINATION ${${PYWRAP}_INSTALL_DIR}
     )
   endif()
@@ -200,23 +217,25 @@ else()
   set(AVX512_COMPILE_OPTION "-mavx512f")
 endif()
 
-CREATE_PYTHON_TARGET(proxsuite_pywrap "" proxsuite)
+create_python_target(proxsuite_pywrap "" proxsuite ${GENERATE_PYTHON_STUBS})
 if(
   BUILD_WITH_VECTORIZATION_SUPPORT
   AND ${CMAKE_SYSTEM_PROCESSOR} MATCHES "(x86)|(X86)|(amd64)|(AMD64)"
 )
   if(BUILD_BINDINGS_WITH_AVX2_SUPPORT)
-    CREATE_PYTHON_TARGET(
+    create_python_target(
       proxsuite_pywrap_avx2
       "${AVX2_COMPILE_OPTION};${FMA_COMPILE_OPTION}"
       proxsuite-vectorized
+      FALSE
     )
   endif(BUILD_BINDINGS_WITH_AVX2_SUPPORT)
   if(BUILD_BINDINGS_WITH_AVX512_SUPPORT)
-    CREATE_PYTHON_TARGET(
+    create_python_target(
       proxsuite_pywrap_avx512
       "${AVX512_COMPILE_OPTION};${FMA_COMPILE_OPTION}"
       proxsuite-vectorized
+      FALSE
     )
   endif(BUILD_BINDINGS_WITH_AVX512_SUPPORT)
 endif()

bindings/python/proxsuite/__init__.py

@@ -1,41 +1,54 @@
-import platform
+"""
+Copyright (c) 2022-2025, INRIA
+"""
+
 import numpy  # noqa F401 for OpenMP proper linkage
+from typing import TYPE_CHECKING
 
-machine = platform.machine()
-has_vectorization_instructions = not machine.startswith(
-    ("arm", "aarch64", "power", "ppc64", "s390x", "sparc")
-)
-if has_vectorization_instructions:
-    from . import instructionset
+if TYPE_CHECKING:
+    from .proxsuite_pywrap import *  # noqa F403
 
 
-def load_main_module(globals):
-    def load_module(main_module_name):
-        import importlib
+def _load_main_module():
+    import platform
+    import importlib
+
+    machine = platform.machine()
+    has_vectorization_instructions = not machine.startswith(
+        ("arm", "aarch64", "power", "ppc64", "s390x", "sparc")
+    )
 
+    def load_module(main_module_name):
         try:
-            main_module = importlib.import_module("." + main_module_name, __name__)
-            globals.update(main_module.__dict__)
-            del globals[main_module_name]
-            return True
+            return importlib.import_module("." + main_module_name, __name__)
         except ModuleNotFoundError:
             return False
 
     if has_vectorization_instructions:  # noqa
+        from . import instructionset
+
         all_modules = [
             ("proxsuite_pywrap_avx512", instructionset.has_AVX512F),
             ("proxsuite_pywrap_avx2", instructionset.has_AVX2),
         ]
 
         for module_name, checker in all_modules:
-            if checker() and load_module(module_name):
-                return
+            if checker() and (mod := load_module(module_name)):
+                return mod
+
+    return load_module("proxsuite_pywrap")
+
+
+_submodule = _load_main_module()
+
 
-    assert load_module("proxsuite_pywrap")
+def __getattr__(name: str):
+    # reroute proxsuite module's attributes to be that of the loaded submodule.
+    return getattr(_submodule, name)
 
 
-load_main_module(globals=globals())
-del load_main_module
-del platform
-del has_vectorization_instructions
-del machine
+def __dir__():
+    # returns iterable of all accessible attributes in the module.
+    # implement this for instropection, for e.g. autocomplete in interpreters (IPython).
+    # We return the attributes from the submodule.
+    return dir(_submodule)