[trajoptlib] Add Python bindings using nanobind

.github/workflows/repair_wheel.py

@@ -0,0 +1,19 @@
+#!/usr/bin/env python3
+
+import os
+import sys
+
+
+def main():
+    for filename in sys.argv[1:]:
+        src = filename
+        dest = src.replace("linux_x86_64", "manylinux_2_35_x86_64").replace(
+            "macosx_14_arm64", "macosx_14_0_arm64"
+        )
+        if src != dest:
+            print(f"{src} -> {dest}")
+            os.rename(src, dest)
+
+
+if __name__ == "__main__":
+    main()

.github/workflows/trajoptlib-py.yml

@@ -0,0 +1,152 @@
+name: Build Python
+
+on: [pull_request, push]
+
+concurrency:
+  group: ${{ github.workflow }}-${{ github.head_ref || github.ref }}
+  cancel-in-progress: true
+
+jobs:
+  build-wheel:
+    timeout-minutes: 10
+    strategy:
+      fail-fast: false
+      matrix:
+        include:
+          - name: Windows x86_64
+            os: windows-2022
+            version: "3.9"
+            cmake-env:
+          - name: Windows x86_64
+            os: windows-2022
+            version: "3.10"
+            cmake-env:
+          - name: Windows x86_64
+            os: windows-2022
+            version: "3.11"
+            cmake-env:
+          - name: Windows x86_64
+            os: windows-2022
+            version: "3.12"
+            cmake-env:
+          - name: Linux x86_64
+            os: ubuntu-24.04
+            version: "3.9"
+            cmake-env:
+          - name: Linux x86_64
+            os: ubuntu-24.04
+            version: "3.10"
+            cmake-env:
+          - name: Linux x86_64
+            os: ubuntu-24.04
+            version: "3.11"
+            cmake-env:
+          - name: Linux x86_64
+            os: ubuntu-24.04
+            version: "3.12"
+            cmake-env:
+          - name: macOS universal
+            os: macOS-14
+            version: "3.10"
+            cmake-env: CMAKE_OSX_ARCHITECTURES="x86_64;arm64"
+          - name: macOS universal
+            os: macOS-14
+            version: "3.11"
+            cmake-env: CMAKE_OSX_ARCHITECTURES="x86_64;arm64"
+          - name: macOS universal
+            os: macOS-14
+            version: "3.12"
+            cmake-env: CMAKE_OSX_ARCHITECTURES="x86_64;arm64"
+
+    name: ${{ matrix.version }} ${{ matrix.name }} wheel
+    runs-on: ${{ matrix.os }}
+    defaults:
+      run:
+        working-directory: ./trajoptlib
+    steps:
+      - uses: actions/checkout@v4
+        with:
+          fetch-depth: 0
+
+      - name: Make GCC 14 the default toolchain (Linux)
+        if: startsWith(matrix.os, 'ubuntu')
+        run: |
+          sudo update-alternatives --install /usr/bin/gcc gcc /usr/bin/gcc-14 200
+          sudo update-alternatives --install /usr/bin/g++ g++ /usr/bin/g++-14 200
+
+      - run: sudo xcode-select -switch /Applications/Xcode_15.3.app
+        if: startsWith(matrix.os, 'macOS')
+
+      - uses: actions/setup-python@v5
+        with:
+          python-version: ${{ matrix.version }}
+
+      - run: pip install typing_extensions
+        if: matrix.version == '3.9' || matrix.version == '3.10'
+
+      - run: pip3 install build pytest
+      - run: ${{ matrix.cmake-env }} python3 -m build --wheel
+      - run: python3 ../.github/workflows/repair_wheel.py trajoptlib-*.whl
+        if: startsWith(matrix.os, 'ubuntu') || startsWith(matrix.os, 'macOS')
+        working-directory: dist
+      - run: pip3 install dist/trajoptlib-*.whl
+        shell: bash
+      - run: pytest
+        shell: bash
+
+      - uses: actions/upload-artifact@v4
+        with:
+          name: ${{ matrix.version }} ${{ matrix.name }} wheel
+          path: dist
+
+  build-sdist:
+    timeout-minutes: 10
+
+    name: sdist
+    runs-on: ubuntu-24.04
+    defaults:
+      run:
+        working-directory: ./trajoptlib
+    steps:
+      - uses: actions/checkout@v4
+        with:
+          fetch-depth: 0
+
+      - uses: actions/setup-python@v5
+        with:
+          python-version: 3.12
+
+      - run: pip3 install build
+      - run: python3 -m build --sdist
+
+      - uses: actions/upload-artifact@v4
+        with:
+          name: sdist
+          path: dist
+
+  pypi-publish:
+    name: Upload release to PyPI
+    runs-on: ubuntu-24.04
+    defaults:
+      run:
+        working-directory: ./trajoptlib
+    needs: [build-wheel, build-sdist]
+    if: github.repository_owner == 'SleipnirGroup' && github.ref == 'refs/heads/main'
+    environment:
+      name: pypi
+      url: https://pypi.org/p/trajoptlib
+    permissions:
+      id-token: write
+    steps:
+      - uses: actions/download-artifact@v4
+        with:
+          path: dist
+          pattern: '* wheel'
+          merge-multiple: true
+      - uses: actions/download-artifact@v4
+        with:
+          path: dist
+          pattern: 'sdist'
+          merge-multiple: true
+      - name: Publish package distributions to PyPI
+        uses: pypa/gh-action-pypi-publish@release/v1

.gitignore

@@ -52,3 +52,6 @@ CMakeUserPresets.json
 
 build/
 cli/
+
+# python
+trajoptlib/.py-build-cmake_cache/

trajoptlib/CMakeLists.txt

@@ -63,6 +63,7 @@ option(BUILD_SHARED_LIBS "Build using shared libraries" ON)
 set(CMAKE_WINDOWS_EXPORT_ALL_SYMBOLS FALSE)
 
 option(BUILD_EXAMPLES "Build examples" OFF)
+option(BUILD_PYTHON "Build python" OFF)
 
 include(CompilerFlags)
 
@@ -83,6 +84,8 @@ target_compile_definitions(TrajoptLib PRIVATE TRAJOPT_EXPORTS)
 include(CTest)
 include(FetchContent)
 
+option(USE_SYSTEM_NANOBIND "Use system nanobind" OFF)
+
 if(BUILD_TESTING)
     # Catch2 dependency
     fetchcontent_declare(
@@ -242,3 +245,89 @@ if(BUILD_EXAMPLES)
         endif()
     endforeach()
 endif()
+
+if(BUILD_PYTHON)
+    find_package(Python REQUIRED COMPONENTS Interpreter Development)
+    if(DEFINED PY_BUILD_CMAKE_MODULE_NAME)
+        set(PY_DEST ${PY_BUILD_CMAKE_MODULE_NAME})
+    else()
+        set(PY_DEST lib/python${Python_VERSION_MAJOR}.${Python_VERSION_MINOR})
+    endif()
+
+    # nanobind dependency
+    if(NOT USE_SYSTEM_NANOBIND)
+        fetchcontent_declare(
+            nanobind
+            GIT_REPOSITORY https://github.com/wjakob/nanobind.git
+            GIT_TAG v2.1.0
+        )
+        fetchcontent_makeavailable(nanobind)
+    else()
+        find_package(nanobind CONFIG REQUIRED)
+    endif()
+
+    file(GLOB_RECURSE trajoptlibpy_src py/cpp/*.cpp)
+
+    # Build Trajoptlib dependency directly into the wheel to avoid having to
+    # configure RPATHs
+    nanobind_add_module(_trajoptlib ${trajoptlibpy_src} ${TrajoptLib_src})
+    compiler_flags(_trajoptlib)
+    target_compile_definitions(_trajoptlib PRIVATE TRAJOPTLIB=1)
+    target_include_directories(
+        _trajoptlib
+        PRIVATE
+            ${CMAKE_CURRENT_SOURCE_DIR}/src
+            ${CMAKE_CURRENT_SOURCE_DIR}/include
+            ${CMAKE_CURRENT_SOURCE_DIR}/py/cpp
+    )
+    target_link_libraries(_trajoptlib PUBLIC Sleipnir)
+
+    # Suppress compiler warnings in nanobind
+    if(${CMAKE_CXX_COMPILER_ID} STREQUAL "AppleClang")
+        target_compile_options(nanobind-static PRIVATE "-Wno-array-bounds")
+    endif()
+
+    install(
+        TARGETS _trajoptlib
+        COMPONENT python_modules
+        LIBRARY
+        DESTINATION ${PY_DEST}
+    )
+
+    nanobind_add_stub(
+        _trajoptlib_stub
+        INSTALL_TIME
+        MARKER_FILE py/py.typed
+        MODULE _trajoptlib
+        OUTPUT trajoptlib/py/__init__.pyi
+        PYTHON_PATH $<TARGET_FILE_DIR:_trajoptlib>
+        DEPENDS _trajoptlib
+        COMPONENT python_modules
+    )
+
+    # # pybind11_mkdoc doesn't support Windows
+    # if(NOT WIN32 AND NOT CMAKE_CROSSCOMPILING)
+    #     # pybind11_mkdoc dependency
+    #     fetchcontent_declare(
+    #         pybind11_mkdoc
+    #         GIT_REPOSITORY https://github.com/pybind/pybind11_mkdoc.git
+    #         GIT_TAG master
+    #         GIT_SUBMODULES ""
+    #     )
+    #     fetchcontent_makeavailable(pybind11_mkdoc)
+    #
+    #     file(
+    #         GLOB_RECURSE trajoptlib_headers
+    #         include/trajopt/*.hpp
+    #         include/trajopt/path/*.hpp
+    #         include/trajopt/geometry/*.hpp
+    #         include/trajopt/constraint/*.hpp
+    #         include/trajopt/util/*.hpp
+    #     )
+    #
+    #     # Generate docs for the Python module
+    #     include(cmake/modules/Pybind11Mkdoc.cmake)
+    #     pybind11_mkdoc(_trajoptlib "${trajoptlib_headers}")
+    #     add_dependencies(_trajoptlib _trajoptlib_docstrings)
+    # endif()
+endif()

trajoptlib/cmake/modules/Pybind11Mkdoc.cmake

@@ -0,0 +1,31 @@
+function(pybind11_mkdoc target headers)
+    find_package(Python3 REQUIRED COMPONENTS Interpreter)
+    if(UNIX AND NOT APPLE)
+        set(env_vars LLVM_DIR_PATH=/usr/lib LIBCLANG_PATH=/usr/lib/libclang.so)
+    endif()
+
+    get_target_property(target_dirs ${target} INCLUDE_DIRECTORIES)
+    list(TRANSFORM target_dirs PREPEND "-I")
+
+    get_target_property(eigen_dirs Eigen3::Eigen INTERFACE_INCLUDE_DIRECTORIES)
+    list(FILTER eigen_dirs INCLUDE REGEX "\\$<BUILD_INTERFACE:.*>")
+    list(TRANSFORM eigen_dirs PREPEND "-I")
+
+    add_custom_command(
+        OUTPUT ${CMAKE_CURRENT_SOURCE_DIR}/jormungandr/cpp/Docstrings.hpp
+        COMMAND
+            ${env_vars} ${Python3_EXECUTABLE} -m pybind11_mkdoc ${headers} -o
+            ${CMAKE_CURRENT_SOURCE_DIR}/jormungandr/cpp/Docstrings.hpp
+            -I/usr/lib/clang/18/include ${target_dirs} ${eigen_dirs} -std=c++23
+        COMMAND
+            ${Python3_EXECUTABLE}
+            ${CMAKE_CURRENT_SOURCE_DIR}/cmake/fix_docstrings.py
+            ${CMAKE_CURRENT_SOURCE_DIR}/jormungandr/cpp/Docstrings.hpp
+        DEPENDS ${headers}
+        USES_TERMINAL
+    )
+    add_custom_target(
+        ${target}_docstrings
+        DEPENDS ${CMAKE_CURRENT_SOURCE_DIR}/jormungandr/cpp/Docstrings.hpp
+    )
+endfunction()

trajoptlib/py/__init__.py

@@ -0,0 +1,5 @@
+"""
+A library to generate time optimal trajectories for FRC robots.
+"""
+
+from ._trajoptlib import *

trajoptlib/py/cpp/BindSwerve.cpp

@@ -0,0 +1,62 @@
+// Copyright (c) TrajoptLib contributors
+
+#include <nanobind/nanobind.h>
+#include <nanobind/stl/function.h>
+
+#include <vector>
+
+#include <sleipnir/optimization/SolverStatus.hpp>
+#include <trajopt/SwerveTrajectoryGenerator.hpp>
+#include <trajopt/geometry/Translation2.hpp>
+
+namespace nb = nanobind;
+
+namespace trajopt {
+/*
+void BindSwerveDrivetrain(nb::class_<SwerveDrivetrain> cls);
+void BindSwerveSolution(nb::class_<SwerveSolution> cls);
+void BindSwerveTrajectorySample(nb::class_<SwerveTrajectorySample> cls);
+void BindSwerveTrajectory(nb::class_<SwerveTrajectory> cls);
+void BindSwervePath(nb::class_<SwervePath> cls);
+void BindSwervePathBuilder(nb::class_<SwervePathBuilder> cls);
+void BindSwerveTrajectoryGenerator(nb::class_<SwerveTrajectoryGenerator> cls);
+ */
+
+void BindSwerveDrivetrain(nb::class_<SwerveDrivetrain>& cls) {
+  using namespace nb::literals;
+
+  cls.def(nb::init<double, double, double, double, double,
+                   std::vector<Translation2d>>(),
+          "mass"_a, "moi"_a, "wheelRadius"_a, "wheelMaxAngularVelocity"_a,
+          "wheelMaxTorque"_a, "modules"_a);
+  cls.def_ro("mass", double, DOC(_trajoptlib, SwerveDrivetrain, double));
+  cls.def_ro("moi", double, DOC(_trajoptlib, SwerveDrivetrain, double));
+  cls.def_ro("wheelRadius", double, DOC(_trajoptlib, SwerveDrivetrain, double));
+  cls.def_ro("wheelMaxAngularVelocity", double, DOC(_trajoptlib, SwerveDrivetrain, double));
+  cls.def_ro("wheelMaxTorque", double, DOC(_trajoptlib, SwerveDrivetrain, double));
+  cls.def_ro("modules", std::vector<Translation2d>, DOC(_trajoptlib, SwerveDrivetrain, std::vector<Translation2d>));
+}
+
+// void BindSwerveSolution(nb::class_<SwerveSolution>& cls) {
+//   using namespace nb::literals;
+//
+//   cls.def(nb::init<
+//               std::vector<double>, std::vector<double>, std::vector<double>, std::vector<double>, std::vector<double>, std::vector<double>, std::vector<double>, std::vector<double>, std::vector<double>, std::vector<double>, std::vector<double>, std::vector<double>, std::vector<double>, >(),
+//           "dt"_a, "x"_a, "y"_a, "thetacos"_a, "thetasin"_a, "vx"_a, "vy"_a,
+//           "omega"_a, "ax"_a, "ay"_a, "alpha"_a, "moduleFX"_a, "moduleFY"_a);
+//   cls.def_ro("dt", std::vector<double>);
+//   cls.def_ro("x", std::vector<double>);
+//   cls.def_ro("y", std::vector<double>);
+//   cls.def_ro("thetacos", std::vector<double>);
+//   cls.def_ro("thetasin", std::vector<double>);
+//   cls.def_ro("vx", std::vector<double>);
+//   cls.def_ro("vy", std::vector<double>);
+//   cls.def_ro("omega", std::vector<double>);
+//   cls.def_ro("ax", std::vector<double>);
+//   cls.def_ro("ay", std::vector<double>);
+//   cls.def_ro("alpha", std::vector<double>);
+//   cls.def_ro("moduleFX", std::vector<double>);
+//   cls.def_ro("moduleFY", std::vector<double>);
+// }
+
+}  // namespace trajopt