From 9eae239059460e793fdb4a29f9ac35ed1b231d22 Mon Sep 17 00:00:00 2001
From: Ryker Fish <rfish@mines.edu>
Date: Tue, 10 Feb 2026 17:38:50 -0700
Subject: [PATCH] feat: add module stubs using nanobind to provide autocomplete
 in python

---
 solvers/CMakeLists.txt              | 24 ++++++++++++++++++++++++
 solvers/DPStokes/CMakeLists.txt     |  8 +++-----
 solvers/NBody/CMakeLists.txt        |  8 +++-----
 solvers/PSE/CMakeLists.txt          |  8 +++-----
 solvers/SelfMobility/CMakeLists.txt |  8 +++-----
 5 files changed, 36 insertions(+), 20 deletions(-)

diff --git a/solvers/CMakeLists.txt b/solvers/CMakeLists.txt
index 8504b66..c9f8fb4 100644
--- a/solvers/CMakeLists.txt
+++ b/solvers/CMakeLists.txt
@@ -1,3 +1,27 @@
+function(setup_nanobind_module NAME)
+  nanobind_add_module(
+    ${NAME}
+    STABLE_ABI
+    python_wrapper.cu
+  )
+
+  nanobind_add_stub(
+    ${NAME}_stub
+    MODULE ${NAME}
+    OUTPUT ${NAME}.pyi
+    DEPENDS ${NAME}
+    MARKER_FILE py.typed
+  )
+
+  install(
+    FILES
+      ${CMAKE_CURRENT_BINARY_DIR}/${NAME}.pyi
+      ${CMAKE_CURRENT_BINARY_DIR}/py.typed
+    DESTINATION ${Python_SITEARCH}/libMobility
+  )
+
+endfunction()
+
 add_subdirectory(DPStokes)
 add_subdirectory(PSE)
 add_subdirectory(NBody)
diff --git a/solvers/DPStokes/CMakeLists.txt b/solvers/DPStokes/CMakeLists.txt
index eaa5952..2420fb1 100644
--- a/solvers/DPStokes/CMakeLists.txt
+++ b/solvers/DPStokes/CMakeLists.txt
@@ -1,11 +1,9 @@
 set(NAME DPStokes)
 add_library(libMobility_${NAME} SHARED extra/uammd_wrapper.cu)
 uammd_setup_target(libMobility_${NAME})
-nanobind_add_module(
-  ${NAME}
-  STABLE_ABI
-  python_wrapper.cu
-)
+
+setup_nanobind_module(${NAME})
+
 uammd_setup_target(${NAME})
 target_link_libraries(${NAME} PRIVATE libMobility_${NAME})
 install(TARGETS libMobility_${NAME} DESTINATION ${CMAKE_INSTALL_PREFIX}/lib)
diff --git a/solvers/NBody/CMakeLists.txt b/solvers/NBody/CMakeLists.txt
index 787587c..3b430df 100644
--- a/solvers/NBody/CMakeLists.txt
+++ b/solvers/NBody/CMakeLists.txt
@@ -1,11 +1,9 @@
 set(NAME NBody)
 add_library(libMobility_${NAME} SHARED extra/NbodyRPY.cu)
 uammd_setup_target(libMobility_${NAME})
-nanobind_add_module(
-  ${NAME}
-  STABLE_ABI
-  python_wrapper.cu
-)
+
+setup_nanobind_module(${NAME})
+
 uammd_setup_target(${NAME})
 target_link_libraries(${NAME} PRIVATE libMobility_${NAME})
 install(TARGETS libMobility_${NAME} DESTINATION ${CMAKE_INSTALL_PREFIX}/lib)
diff --git a/solvers/PSE/CMakeLists.txt b/solvers/PSE/CMakeLists.txt
index 7237403..1ba6093 100644
--- a/solvers/PSE/CMakeLists.txt
+++ b/solvers/PSE/CMakeLists.txt
@@ -1,11 +1,9 @@
 set(NAME PSE)
 add_library(libMobility_${NAME} SHARED extra/uammd_wrapper.cu)
 uammd_setup_target(libMobility_${NAME})
-nanobind_add_module(
-  ${NAME}
-  STABLE_ABI
-  python_wrapper.cu
-)
+
+setup_nanobind_module(${NAME})
+
 uammd_setup_target(${NAME})
 target_link_libraries(${NAME} PRIVATE libMobility_${NAME})
 install(TARGETS libMobility_${NAME} DESTINATION ${CMAKE_INSTALL_PREFIX}/lib)
diff --git a/solvers/SelfMobility/CMakeLists.txt b/solvers/SelfMobility/CMakeLists.txt
index 50e1247..022f440 100644
--- a/solvers/SelfMobility/CMakeLists.txt
+++ b/solvers/SelfMobility/CMakeLists.txt
@@ -1,11 +1,9 @@
 set(NAME SelfMobility)
 add_library(libMobility_${NAME} SHARED selfmobility.cu)
 uammd_setup_target(libMobility_${NAME})
-nanobind_add_module(
-  ${NAME}
-  STABLE_ABI
-  python_wrapper.cu
-)
+
+setup_nanobind_module(${NAME})
+
 uammd_setup_target(${NAME})
 target_link_libraries(${NAME} PRIVATE libMobility_${NAME})
 install(TARGETS libMobility_${NAME} DESTINATION ${CMAKE_INSTALL_PREFIX}/lib)