sleipnirgroup-jormungandr 0.6.7.dev56__tar.gz → 0.6.7.dev59__tar.gz

{sleipnirgroup_jormungandr-0.6.7.dev56 → sleipnirgroup_jormungandr-0.6.7.dev59}/CMakeLists.txt

@@ -301,7 +301,9 @@ if(SLEIPNIR_BUILD_PYTHON)
     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})
+        set(PY_DEST
+            lib/python${Python_VERSION_MAJOR}.${Python_VERSION_MINOR}/sleipnir
+        )
     endif()
 
     # nanobind dependency
@@ -343,47 +345,57 @@ if(SLEIPNIR_BUILD_PYTHON)
         LIBRARY DESTINATION ${PY_DEST}
     )
 
+    # Generate stubs at install time because only the install tree has the
+    # sleipnir/ package layout needed to import the module by its fully
+    # qualified name `sleipnir._sleipnir`.
+    #
+    # Importing sleipnir._sleipnir requires the directory containing the
+    # sleipnir package (i.e., PY_DEST's parent) to be on the Python path. For
+    # py-build-cmake builds, PY_DEST has no parent. For plain CMake installs,
+    # PY_DEST's parent is lib/pythonX.Y.
+    cmake_path(GET PY_DEST PARENT_PATH PY_DEST_PARENT)
     nanobind_add_stub(
         _sleipnir_stub
         INSTALL_TIME
-        MARKER_FILE sleipnir/py.typed
-        MODULE _sleipnir
-        OUTPUT sleipnir/__init__.pyi
-        PYTHON_PATH $<TARGET_FILE_DIR:_sleipnir>
+        MARKER_FILE ${PY_DEST}/py.typed
+        MODULE sleipnir._sleipnir
+        OUTPUT ${PY_DEST}/_sleipnir/__init__.pyi
+        PYTHON_PATH ${PY_DEST_PARENT}
         DEPENDS _sleipnir
         COMPONENT python_modules
     )
     nanobind_add_stub(
         _sleipnir_autodiff_stub
         INSTALL_TIME
-        MODULE _sleipnir.autodiff
-        OUTPUT sleipnir/autodiff/__init__.pyi
-        PYTHON_PATH $<TARGET_FILE_DIR:_sleipnir>
+        MODULE sleipnir._sleipnir.autodiff
+        OUTPUT ${PY_DEST}/_sleipnir/autodiff.pyi
+        PYTHON_PATH ${PY_DEST_PARENT}
         DEPENDS _sleipnir
         COMPONENT python_modules
     )
     nanobind_add_stub(
         _sleipnir_optimization_stub
         INSTALL_TIME
-        MODULE _sleipnir.optimization
-        OUTPUT sleipnir/optimization/__init__.pyi
-        PYTHON_PATH $<TARGET_FILE_DIR:_sleipnir>
+        MODULE sleipnir._sleipnir.optimization
+        OUTPUT ${PY_DEST}/_sleipnir/optimization.pyi
+        PYTHON_PATH ${PY_DEST_PARENT}
         DEPENDS _sleipnir
         COMPONENT python_modules
     )
+    install(
+        CODE
+            "execute_process(
+                COMMAND \"${Python_EXECUTABLE}\"
+                    \"${CMAKE_CURRENT_SOURCE_DIR}/python/postprocess_stubs.py\"
+                    ${PY_DEST}/_sleipnir/autodiff.pyi
+                    ${PY_DEST}/_sleipnir/optimization.pyi
+                WORKING_DIRECTORY \"\${CMAKE_INSTALL_PREFIX}\"
+                COMMAND_ERROR_IS_FATAL ANY
+            )"
+        COMPONENT python_modules
+    )
 
     if(NOT CMAKE_CROSSCOMPILING)
-        # pybind11_mkdoc dependency
-        FetchContent_Declare(
-            pybind11_mkdoc
-            GIT_REPOSITORY https://github.com/pybind/pybind11_mkdoc.git
-            GIT_TAG v3.0.0
-            GIT_SHALLOW ON
-            EXCLUDE_FROM_ALL
-            SYSTEM
-        )
-        FetchContent_MakeAvailable(pybind11_mkdoc)
-
         set(sleipnir_headers
             ${CMAKE_CURRENT_SOURCE_DIR}/include/sleipnir/autodiff/expression_type.hpp
             ${CMAKE_CURRENT_SOURCE_DIR}/include/sleipnir/autodiff/gradient.hpp

{sleipnirgroup_jormungandr-0.6.7.dev56 → sleipnirgroup_jormungandr-0.6.7.dev59}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: sleipnirgroup-jormungandr
-Version: 0.6.7.dev56
+Version: 0.6.7.dev59
 Summary: Reverse mode autodiff library and NLP solver DSL
 License-Expression: BSD-3-Clause
 License-File: LICENSE.txt

{sleipnirgroup_jormungandr-0.6.7.dev56 → sleipnirgroup_jormungandr-0.6.7.dev59}/cmake/Pybind11Mkdoc.cmake

@@ -1,6 +1,20 @@
 function(pybind11_mkdoc target headers)
     find_package(Python3 REQUIRED COMPONENTS Interpreter)
 
+    # Check that pybind11_mkdoc is installed
+    execute_process(
+        COMMAND ${Python3_EXECUTABLE} -c "import pybind11_mkdoc"
+        RESULT_VARIABLE pybind11_mkdoc_result
+        OUTPUT_QUIET
+        ERROR_QUIET
+    )
+    if(NOT pybind11_mkdoc_result EQUAL 0)
+        message(
+            FATAL_ERROR
+            "pybind11_mkdoc not found for ${Python3_EXECUTABLE}. Install it with `pip install --group docstrings`."
+        )
+    endif()
+
     # Target compiler flags
     get_target_property(target_flags ${target} INCLUDE_DIRECTORIES)
     list(TRANSFORM target_flags PREPEND "-I")

{sleipnirgroup_jormungandr-0.6.7.dev56 → sleipnirgroup_jormungandr-0.6.7.dev59}/pyproject.toml

@@ -21,8 +21,11 @@ build-backend = "py_build_cmake.build"
 
 [dependency-groups]
 dev = [
+    "mypy>=2.4.0",
     "pytest>=9.1.1",
+    "scipy-stubs",
 ]
+docstrings = ["pybind11-mkdoc~=3.0.0"]
 
 [tool.py-build-cmake.cmake]
 build_type = "Release"
@@ -54,8 +57,27 @@ exclude = []
 python_abi = "abi3"
 abi3_minimum_cpython_version = 312
 
+[tool.mypy]
+files = [
+  "benchmarks",
+  "docs/generate_website.py",
+  "examples",
+  "python/postprocess_stubs.py",
+  "python/test",
+  "tools",
+]
+allow_redefinition_new = true
+explicit_package_bases = true
+local_partial_types = true
+mypy_path = ["examples/current_manager", "python/test/optimization"]
+
+  [[tool.mypy.overrides]]
+  module = "mpl_toolkits.*"
+  ignore_missing_imports = true
+
 [tool.pytest.ini_options]
 minversion = "6.0"
+pythonpath = ["examples/current_manager"]
 testpaths = ["python/test", "examples/current_manager/test"]
 
 [tool.uv]
@@ -68,6 +90,7 @@ cache-keys = [
     { file = "python/cpp/**/*.cpp" },
     { file = "python/cpp/binders.hpp" },
     { file = "python/cpp/for_each_type.hpp" },
+    { file = "python/cpp/slice_tuple.hpp" },
     { file = "python/cpp/try_cast.hpp" },
     { file = "python/src/**" },
     { file = "src/**" },

{sleipnirgroup_jormungandr-0.6.7.dev56 → sleipnirgroup_jormungandr-0.6.7.dev59}/python/cpp/autodiff/bind_variable_block.cpp

@@ -6,6 +6,7 @@
 #include <nanobind/operators.h>
 #include <nanobind/stl/function.h>
 #include <nanobind/stl/string.h>
+#include <nanobind/stl/tuple.h>
 #include <nanobind/stl/vector.h>
 
 #include <stdint.h>
@@ -15,6 +16,7 @@
 #include <functional>
 #include <stdexcept>
 #include <string>
+#include <tuple>
 
 #include <Eigen/Core>
 #include <sleipnir/autodiff/slice.hpp>
@@ -25,6 +27,7 @@
 
 #include "docstrings.hpp"
 #include "for_each_type.hpp"
+#include "slice_tuple.hpp"
 #include "try_cast.hpp"
 
 namespace nb = nanobind;
@@ -68,7 +71,7 @@ void bind_variable_block(
       "row"_a, "value"_a);
   cls.def(
       "__setitem__",
-      [](VariableBlock<VariableMatrix<double>>& self, nb::tuple slices,
+      [](VariableBlock<VariableMatrix<double>>& self, SliceTuple slices,
          nb::object value) {
         if (slices.size() != 2) {
           throw nb::index_error(
@@ -151,31 +154,32 @@ void bind_variable_block(
   cls.def(
       "__getitem__",
       [](VariableBlock<VariableMatrix<double>>& self,
-         nb::tuple slices) -> nb::object {
+         std::tuple<int, int> slices) -> Variable<double>& {
+        auto [row, col] = slices;
+
+        if (row >= self.rows() || col >= self.cols()) {
+          throw std::out_of_range("Index out of bounds");
+        }
+
+        if (row < 0) {
+          row += self.rows();
+        }
+        if (col < 0) {
+          col += self.cols();
+        }
+        return self[row, col];
+      },
+      nb::rv_policy::reference_internal, "slices"_a,
+      DOC(slp, VariableBlock, operator, array));
+  cls.def(
+      "__getitem__",
+      [](VariableBlock<VariableMatrix<double>>& self,
+         SliceTuple slices) -> VariableBlock<VariableMatrix<double>> {
         if (slices.size() != 2) {
           throw nb::index_error(
               std::format("Expected 2 slices, got {}.", slices.size()).c_str());
         }
 
-        // If both indices are integers instead of slices, return Variable
-        // instead of VariableBlock
-        if (nb::isinstance<int>(slices[0]) && nb::isinstance<int>(slices[1])) {
-          int row = nb::cast<int>(slices[0]);
-          int col = nb::cast<int>(slices[1]);
-
-          if (row >= self.rows() || col >= self.cols()) {
-            throw std::out_of_range("Index out of bounds");
-          }
-
-          if (row < 0) {
-            row += self.rows();
-          }
-          if (col < 0) {
-            col += self.cols();
-          }
-          return nb::cast(self[row, col]);
-        }
-
         Slice row_slice;
         int row_slice_length;
         Slice col_slice;
@@ -211,8 +215,7 @@ void bind_variable_block(
           col_slice_length = 1;
         }
 
-        return nb::cast(
-            self[row_slice, row_slice_length, col_slice, col_slice_length]);
+        return self[row_slice, row_slice_length, col_slice, col_slice_length];
       },
       nb::keep_alive<0, 1>(), DOC(slp, VariableBlock, operator, array),
       "slices"_a);
@@ -223,6 +226,31 @@ void bind_variable_block(
           nb::overload_cast<int>(&VariableBlock<VariableMatrix<double>>::col),
           "col"_a, DOC(slp, VariableBlock, col));
 
+  // Matrix-matrix multiplication. Products with a VariableMatrix on the left
+  // are bound on VariableMatrix.
+  for (const char* name : {"__mul__", "__matmul__"}) {
+    cls.def(
+        name,
+        [](const VariableBlock<VariableMatrix<double>>& lhs,
+           const VariableBlock<VariableMatrix<double>>& rhs) {
+          return lhs * rhs;
+        },
+        nb::is_operator(), "rhs"_a);
+  }
+
+  // Matrix multiplication with NumPy arrays. These are also handled by
+  // __array_ufunc__ below, but type checkers only see these overloads.
+  cls.def(
+      "__matmul__",
+      [](const VariableBlock<VariableMatrix<double>>& lhs,
+         nb::DRef<Eigen::MatrixXd> rhs) { return lhs * rhs; },
+      nb::is_operator(), "rhs"_a);
+  cls.def(
+      "__rmatmul__",
+      [](const VariableBlock<VariableMatrix<double>>& rhs,
+         nb::DRef<Eigen::MatrixXd> lhs) { return lhs * rhs; },
+      nb::is_operator(), "lhs"_a);
+
   // https://numpy.org/doc/stable/user/basics.dispatch.html
   cls.def(
       "__array_ufunc__",

{sleipnirgroup_jormungandr-0.6.7.dev56 → sleipnirgroup_jormungandr-0.6.7.dev59}/python/cpp/autodiff/bind_variable_matrix.cpp

@@ -6,6 +6,7 @@
 #include <nanobind/operators.h>
 #include <nanobind/stl/function.h>
 #include <nanobind/stl/string.h>
+#include <nanobind/stl/tuple.h>
 #include <nanobind/stl/vector.h>
 
 #include <stdint.h>
@@ -15,6 +16,7 @@
 #include <functional>
 #include <stdexcept>
 #include <string>
+#include <tuple>
 #include <vector>
 
 #include <Eigen/Core>
@@ -26,6 +28,7 @@
 
 #include "docstrings.hpp"
 #include "for_each_type.hpp"
+#include "slice_tuple.hpp"
 #include "try_cast.hpp"
 
 namespace nb = nanobind;
@@ -90,7 +93,7 @@ void bind_variable_matrix(nb::module_& autodiff,
       "row"_a, "value"_a);
   cls.def(
       "__setitem__",
-      [](VariableMatrix<double>& self, nb::tuple slices, nb::object value) {
+      [](VariableMatrix<double>& self, SliceTuple slices, nb::object value) {
         if (slices.size() != 2) {
           throw nb::index_error(
               std::format("Expected 2 slices, got {}.", slices.size()).c_str());
@@ -170,31 +173,33 @@ void bind_variable_matrix(nb::module_& autodiff,
       DOC(slp, VariableMatrix, operator, array, 3));
   cls.def(
       "__getitem__",
-      [](VariableMatrix<double>& self, nb::tuple slices) -> nb::object {
+      [](VariableMatrix<double>& self,
+         std::tuple<int, int> slices) -> Variable<double>& {
+        auto [row, col] = slices;
+
+        if (row >= self.rows() || col >= self.cols()) {
+          throw std::out_of_range("Index out of bounds");
+        }
+
+        if (row < 0) {
+          row += self.rows();
+        }
+        if (col < 0) {
+          col += self.cols();
+        }
+        return self[row, col];
+      },
+      nb::rv_policy::reference_internal, "slices"_a,
+      DOC(slp, VariableMatrix, operator, array));
+  cls.def(
+      "__getitem__",
+      [](VariableMatrix<double>& self,
+         SliceTuple slices) -> VariableBlock<VariableMatrix<double>> {
         if (slices.size() != 2) {
           throw nb::index_error(
               std::format("Expected 2 slices, got {}.", slices.size()).c_str());
         }
 
-        // If both indices are integers instead of slices, return Variable
-        // instead of VariableBlock
-        if (nb::isinstance<int>(slices[0]) && nb::isinstance<int>(slices[1])) {
-          int row = nb::cast<int>(slices[0]);
-          int col = nb::cast<int>(slices[1]);
-
-          if (row >= self.rows() || col >= self.cols()) {
-            throw std::out_of_range("Index out of bounds");
-          }
-
-          if (row < 0) {
-            row += self.rows();
-          }
-          if (col < 0) {
-            col += self.cols();
-          }
-          return nb::cast(self[row, col]);
-        }
-
         Slice row_slice;
         int row_slice_length;
         Slice col_slice;
@@ -227,8 +232,7 @@ void bind_variable_matrix(nb::module_& autodiff,
           col_slice_length = 1;
         }
 
-        return nb::cast(
-            self[row_slice, row_slice_length, col_slice, col_slice_length]);
+        return self[row_slice, row_slice_length, col_slice, col_slice_length];
       },
       nb::keep_alive<0, 1>(), "slices"_a,
       DOC(slp, VariableMatrix, operator, array));
@@ -237,43 +241,46 @@ void bind_variable_matrix(nb::module_& autodiff,
   cls.def("col", nb::overload_cast<int>(&VariableMatrix<double>::col), "col"_a,
           DOC(slp, VariableMatrix, col));
 
-  // Matrix-matrix multiplication
-  cls.def(
-      "__mul__",
-      [](const VariableMatrix<double>& lhs, const VariableMatrix<double>& rhs) {
-        return lhs * rhs;
-      },
-      nb::is_operator(), "rhs"_a);
+  // Matrix-matrix multiplication. Products with a VariableBlock on the left are
+  // bound on VariableBlock.
+  for (const char* name : {"__mul__", "__matmul__"}) {
+    cls.def(
+        name,
+        [](const VariableMatrix<double>& lhs,
+           const VariableMatrix<double>& rhs) { return lhs * rhs; },
+        nb::is_operator(), "rhs"_a);
+    cls.def(
+        name,
+        [](const VariableMatrix<double>& lhs,
+           const VariableBlock<VariableMatrix<double>>& rhs) {
+          return lhs * rhs;
+        },
+        nb::is_operator(), "rhs"_a);
+  }
+  for (const char* name : {"__rmul__", "__rmatmul__"}) {
+    cls.def(
+        name,
+        [](const VariableMatrix<double>& rhs,
+           const VariableBlock<VariableMatrix<double>>& lhs) {
+          return lhs * rhs;
+        },
+        nb::is_operator(), "lhs"_a);
+  }
+
+  // Matrix multiplication with NumPy arrays. These are also handled by
+  // __array_ufunc__ below, but type checkers only see these overloads.
   cls.def(
-      "__mul__",
-      [](const VariableMatrix<double>& lhs,
-         const VariableBlock<VariableMatrix<double>>& rhs) {
+      "__matmul__",
+      [](const VariableMatrix<double>& lhs, nb::DRef<Eigen::MatrixXd> rhs) {
         return lhs * rhs;
       },
       nb::is_operator(), "rhs"_a);
   cls.def(
-      "__rmul__",
-      [](const VariableMatrix<double>& rhs,
-         const VariableBlock<VariableMatrix<double>>& lhs) {
+      "__rmatmul__",
+      [](const VariableMatrix<double>& rhs, nb::DRef<Eigen::MatrixXd> lhs) {
         return lhs * rhs;
       },
       nb::is_operator(), "lhs"_a);
-  cls.def(
-      "__mul__",
-      [](const VariableBlock<VariableMatrix<double>>& lhs,
-         const VariableBlock<VariableMatrix<double>>& rhs) {
-        return lhs * rhs;
-      },
-      nb::is_operator(), "rhs"_a);
-  for_each_type<VariableMatrix<double>, VariableBlock<VariableMatrix<double>>>(
-      [&]<typename L> {
-        for_each_type<VariableMatrix<double>,
-                      VariableBlock<VariableMatrix<double>>>([&]<typename R> {
-          cls.def(
-              "__matmul__", [](L&& lhs, R&& rhs) { return lhs * rhs; },
-              nb::is_operator(), "rhs"_a);
-        });
-      });
 
   // https://numpy.org/doc/stable/user/basics.dispatch.html
   cls.def(

{sleipnirgroup_jormungandr-0.6.7.dev56 → sleipnirgroup_jormungandr-0.6.7.dev59}/python/cpp/optimization/bind_ocp.cpp

@@ -21,6 +21,23 @@
 
 namespace nb = nanobind;
 
+namespace nanobind::detail {
+
+// nanobind's duration caster accepts a float in seconds as well as a
+// timedelta, but its type annotation only advertises timedelta because float is
+// an implicit conversion. Advertise float in the stub's input type so mypy
+// accepts passing a float to the dt parameter of OCP's constructor,
+// set_min_timestep(), and set_max_timestep().
+template <>
+class type_caster<std::chrono::duration<double>>
+    : public duration_caster<std::chrono::duration<double>> {
+ public:
+  static constexpr auto Name =
+      io_name("datetime.timedelta | float", "datetime.timedelta");
+};
+
+}  // namespace nanobind::detail
+
 namespace slp {
 
 void bind_ocp(nb::class_<OCP<double>, Problem<double>>& cls) {

sleipnirgroup_jormungandr-0.6.7.dev59/python/cpp/slice_tuple.hpp

@@ -0,0 +1,21 @@
+// Copyright (c) Sleipnir contributors
+
+#pragma once
+
+#include <nanobind/nanobind.h>
+#include <nanobind/stl/variant.h>
+
+#include <variant>
+
+namespace nb = nanobind;
+
+namespace slp {
+
+/// A 2D index of the form (row, col) where each element is an int or a slice.
+///
+/// The element types only affect the generated type stubs; the tuple's contents
+/// are checked at runtime.
+using SliceTuple = nb::typed<nb::tuple, std::variant<int, nb::slice>,
+                             std::variant<int, nb::slice>>;
+
+}  // namespace slp

sleipnirgroup_jormungandr-0.6.7.dev59/python/postprocess_stubs.py

@@ -0,0 +1,166 @@
+#!/usr/bin/env python
+
+"""
+Fixes up type stubs generated by nanobind's stubgen.
+
+stubgen can't express some things the bindings support at runtime:
+
+1. VariableBlock implicitly converts to VariableMatrix, so every VariableMatrix
+   parameter also accepts a VariableBlock. Return types and callback arguments
+   are left alone since those are always VariableMatrix.
+2. Sparse matrices are annotated as scipy.sparse.csc_matrix[float], but
+   scipy-stubs expects a NumPy numeric scalar type argument.
+"""
+
+import argparse
+import ast
+import re
+from pathlib import Path
+
+PREFIX = r"(?:sleipnir\._sleipnir\.autodiff\.)?"
+
+# VariableMatrix, optionally module-qualified. The lookahead prevents
+# already-widened annotations from matching again.
+#
+#   x: VariableMatrix | None
+#   -> x: VariableMatrix | VariableBlock | None
+VARIABLE_MATRIX = re.compile(
+    rf"""
+    (?<![\w.])                                  # Not part of another name
+    ({PREFIX})                                  # Optional module prefix
+    VariableMatrix\b                            # Type name
+    (?!\ \|\ {PREFIX}VariableBlock\b)           # Not already widened
+    """,
+    re.VERBOSE,
+)
+
+
+def is_callable(annotation: ast.expr) -> bool:
+    """Returns True if an annotation contains a Callable type."""
+    return any(
+        (isinstance(node, ast.Name) and node.id == "Callable")
+        or (isinstance(node, ast.Attribute) and node.attr == "Callable")
+        for node in ast.walk(annotation)
+    )
+
+
+def parameters(node: ast.FunctionDef) -> list[ast.arg | None]:
+    """Returns a function's parameters in signature order."""
+    args = node.args
+    return [
+        *args.posonlyargs,
+        None,  # Separates positional-only parameters from the rest
+        *args.args,
+        args.vararg,
+        *args.kwonlyargs,
+        args.kwarg,
+    ]
+
+
+def signature(node: ast.FunctionDef) -> list[str | None]:
+    """Returns a function's parameter annotations as source strings."""
+    return [
+        None if arg is None or arg.annotation is None else ast.unparse(arg.annotation)
+        for arg in parameters(node)
+    ]
+
+
+def widen_signatures(stub: str) -> str:
+    """
+    Widens every VariableMatrix in function parameter annotations to
+    VariableMatrix | VariableBlock.
+
+    A parameter isn't widened if a sibling overload is identical except it
+    takes VariableBlock there; widening it would make that overload
+    unreachable.
+    """
+    # Group overloads by name within each module or class body
+    overload_sets: dict[tuple[int, str], list[ast.FunctionDef]] = {}
+    for node in ast.walk(ast.parse(stub)):
+        for child in getattr(node, "body", []):
+            if isinstance(child, ast.FunctionDef):
+                overload_sets.setdefault((id(node), child.name), []).append(child)
+
+    annotations = []
+    for overloads in overload_sets.values():
+        signatures = [signature(overload) for overload in overloads]
+        for overload, sig in zip(overloads, signatures):
+            for i, arg in enumerate(parameters(overload)):
+                # Callbacks receive VariableMatrix, so their annotations are
+                # left alone
+                if arg is None or arg.annotation is None or is_callable(arg.annotation):
+                    continue
+
+                # Skip if a sibling overload already takes VariableBlock here
+                block_sig = sig.copy()
+                block_sig[i] = VARIABLE_MATRIX.sub(
+                    r"\1VariableBlock", ast.unparse(arg.annotation)
+                )
+                if block_sig != sig and block_sig in signatures:
+                    continue
+
+                annotations.append(arg.annotation)
+
+    # Encode stub string into UTF-8 bytes
+    source_bytes = stub.encode()
+
+    # Compute byte offsets of each source line
+    line_byte_offsets = [0]
+    for line in source_bytes.splitlines(keepends=True):
+        line_byte_offsets.append(line_byte_offsets[-1] + len(line))
+
+    # Edit back to front so earlier offsets stay valid
+    for annotation in sorted(
+        annotations, key=lambda node: (node.lineno, node.col_offset), reverse=True
+    ):
+        # Assertions to satisfy mypy. They can't fail in practice because
+        # ast.parse() always sets those fields.
+        assert annotation.end_lineno is not None
+        assert annotation.end_col_offset is not None
+
+        # Compute annotation's byte offsets relative to start of file
+        start = line_byte_offsets[annotation.lineno - 1] + annotation.col_offset
+        end = line_byte_offsets[annotation.end_lineno - 1] + annotation.end_col_offset
+
+        # Splice widened type annotations into source_bytes
+        #
+        # \1 is "sleipnir._sleipnir.autodiff." prefix or empty
+        widened = VARIABLE_MATRIX.sub(
+            r"\1VariableMatrix | \1VariableBlock", source_bytes[start:end].decode()
+        )
+        source_bytes = source_bytes[:start] + widened.encode() + source_bytes[end:]
+    return source_bytes.decode()
+
+
+def fix_sparse_matrices(stub: str) -> str:
+    """
+    Replaces scipy.sparse.csc_matrix[float] with
+    scipy.sparse.csc_matrix[numpy.float64], since scipy-stubs requires a NumPy
+    numeric scalar type argument.
+    """
+    return stub.replace(
+        "scipy.sparse.csc_matrix[float]", "scipy.sparse.csc_matrix[numpy.float64]"
+    )
+
+
+def main() -> None:
+    parser = argparse.ArgumentParser(
+        description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter
+    )
+    parser.add_argument(
+        "stubs",
+        nargs="+",
+        type=Path,
+        help="stub files (.pyi) to fix up in place",
+    )
+    args = parser.parse_args()
+
+    for path in args.stubs:
+        stub = path.read_text(encoding="utf-8")
+        stub = widen_signatures(stub)
+        stub = fix_sparse_matrices(stub)
+        path.write_text(stub, encoding="utf-8")
+
+
+if __name__ == "__main__":
+    main()

{sleipnirgroup_jormungandr-0.6.7.dev56 → sleipnirgroup_jormungandr-0.6.7.dev59}/python/src/sleipnir/__init__.py

@@ -2,4 +2,4 @@
 
 from ._sleipnir import *
 
-__version__ = "0.6.7.dev56"
+__version__ = "0.6.7.dev59"