Switch to nanobind

.gitignore

@@ -20,3 +20,6 @@ libmessaging.*
 libmessaging_shared.*
 .sconsign.dblite
 .mypy_cache/
+
+/msgq/ipc_pyx.pyi
+/msgq/visionipc/visionipc_pyx.pyi

MANIFEST.in

@@ -1,4 +1,3 @@
 include README.md pyproject.toml setup.py SConstruct SConscript
-recursive-include site_scons *.py
-recursive-include msgq *.py *.pyx *.pxd *.h *.cc
-global-exclude *.so *.o *.os *.a *.pyc
+recursive-include msgq *.py py.typed *.h *.cc
+global-exclude *.pyi *.so *.o *.os *.a *.pyc

README.md

@@ -107,6 +107,8 @@ msgq.MultiplePublishersError                     # Publisher conflict; subclass
 
 Issues and pull requests are welcome on [GitHub](https://github.com/commaai/msgq). Run `./test.sh` to build, lint, and test the package.
 
+The Python bindings use nanobind. SCons generates type stubs from the compiled modules for type checking and includes them in the Python package.
+
 ## License
 
 MSGQ is available under the [MIT License](https://github.com/commaai/msgq/blob/master/LICENSE).

SConscript

@@ -1,7 +1,11 @@
-Import('env', 'envCython', 'common')
+import sys
+import sysconfig
+
+Import('env', 'envNanobind', 'nanobind_lib', 'common')
 
 
 visionipc_dir = Dir('msgq/visionipc')
+python_suffix = sysconfig.get_config_var('EXT_SUFFIX')
 
 # Build msgq
 msgq_objects = env.SharedObject([
@@ -10,7 +14,8 @@ msgq_objects = env.SharedObject([
   'msgq/msgq.cc',
 ])
 msgq = env.Library('msgq', msgq_objects)
-msgq_python = envCython.Program('msgq/ipc_pyx.so', 'msgq/ipc_pyx.pyx', LIBS=envCython["LIBS"]+[msgq]+common)
+# Keep the extension module names compatible with existing Python imports.
+msgq_python = envNanobind.Program(f'msgq/ipc_pyx{python_suffix}', ['msgq/ipc_py.cc', nanobind_lib], LIBS=envNanobind["LIBS"]+[msgq]+common)
 
 # Build Vision IPC
 vipc_files = ['visionipc.cc', 'visionipc_server.cc', 'visionipc_client.cc']
@@ -24,10 +29,16 @@ vipc_objects = env.SharedObject(vipc_sources)
 visionipc = env.Library('visionipc', vipc_objects)
 
 
-vipc_libs = envCython["LIBS"] + [visionipc, msgq] + common
-envCython.Program(f'{visionipc_dir.abspath}/visionipc_pyx.so', f'{visionipc_dir.abspath}/visionipc_pyx.pyx',
+vipc_libs = envNanobind["LIBS"] + [visionipc, msgq] + common
+visionipc_python = envNanobind.Program(f'{visionipc_dir.abspath}/visionipc_pyx{python_suffix}', [f'{visionipc_dir.abspath}/visionipc_py.cc', nanobind_lib],
                   LIBS=vipc_libs)
 
+# ASan modules require a preloaded runtime before Python can import them.
+if not GetOption('asan'):
+  for module in ('msgq.ipc_pyx', 'msgq.visionipc.visionipc_pyx'):
+    env.Command(module.replace('.', '/') + '.pyi', [msgq_python, visionipc_python],
+                f'"{sys.executable}" -m nanobind.stubgen -i "{Dir("#").abspath}" -m {module} -o "$TARGET"')
+
 if GetOption('extras'):
   env.Program('msgq/test_runner', ['msgq/msgq_tests.cc'], LIBS=[msgq]+common)
 

SConstruct

@@ -1,8 +1,11 @@
 import os
 import platform
+from pathlib import Path
 import subprocess
 import sysconfig
 
+import nanobind
+
 arch = subprocess.check_output(["uname", "-m"], encoding='utf8').rstrip()
 if platform.system() == "Darwin":
   arch = "Darwin"
@@ -62,21 +65,24 @@ env = Environment(
   CFLAGS="-std=gnu11",
   CXXFLAGS="-std=c++1z",
   CPPPATH=cpppath,
-  CYTHONCFILESUFFIX=".cpp",
-  tools=["default", "cython"]
+  tools=["default"]
 )
 
 Export('env', 'arch', 'common')
 
-envCython = env.Clone(LIBS=[])
-envCython["CCFLAGS"] += ["-Wno-#warnings", "-Wno-cpp", "-Wno-shadow", "-Wno-deprecated-declarations"]
-envCython["CCFLAGS"].remove('-Werror')
+envNanobind = env.Clone(LIBS=[])
+nb_dir = Path(nanobind.source_dir()).parent
+envNanobind.Append(CPPPATH=[nanobind.include_dir(), str(nb_dir / "ext/robin_map/include")],
+                   CCFLAGS=["-fvisibility=hidden", "-fno-strict-aliasing"])
 if arch == "Darwin":
-  envCython["LINKFLAGS"] = ["-bundle", "-undefined", "dynamic_lookup"]
+  envNanobind["LINKFLAGS"] = ldflags + ["-bundle", "-undefined", "dynamic_lookup"]
 else:
-  envCython["LINKFLAGS"] = ["-pthread", "-shared"]
-
-Export('envCython')
+  envNanobind["LINKFLAGS"] = ldflags + ["-pthread", "-shared"]
 
+# Third-party runtime sources do not follow our shadowing warning policy.
+envNanobindRuntime = envNanobind.Clone()
+envNanobindRuntime['CCFLAGS'] = [flag for flag in envNanobindRuntime['CCFLAGS'] if not flag.startswith('-Wshadow')]
+nanobind_lib = envNanobindRuntime.SharedObject('#/build/nb_combined.os', str(nb_dir / "src/nb_combined.cpp"))
+Export('envNanobind', 'nanobind_lib')
 
 SConscript(['SConscript'])

lefthook.yml

@@ -18,12 +18,12 @@ test:
     ty:
       run: ty check .
     codespell:
-      run: codespell {files} -L ned,stdio,master --builtin clear,rare,informal,usage,code,names,en-GB_to_en-US -S uv.lock,*_pyx.cpp
+      run: codespell {files} -L ned,stdio,master --builtin clear,rare,informal,usage,code,names,en-GB_to_en-US -S uv.lock
       files: git ls-tree -r HEAD --name-only
     cppcheck:
       run: cppcheck --error-exitcode=1 --inline-suppr --language=c++ --force --quiet -j4 --check-level=exhaustive $(git ls-files '*.cc' | grep -v -E '(msgq_tests|test_runner)\.cc')
     cpplint:
-      run: cpplint --exclude=msgq/ipc_pyx.cpp --exclude=msgq/visionipc/visionipc_pyx.cpp --recursive --quiet --counting=detailed --linelength=240 --filter=-build,-legal,-readability,-runtime,-whitespace,+build/include_subdir,+build/forward_decl,+build/include_what_you_use,+build/deprecated,+whitespace/comma,+whitespace/line_length,+whitespace/empty_if_body,+whitespace/empty_loop_body,+whitespace/empty_conditional_body,+whitespace/forcolon,+whitespace/parens,+whitespace/semicolon,+whitespace/tab,+readability/braces msgq/
+      run: cpplint --recursive --quiet --counting=detailed --linelength=240 --filter=-build,-legal,-readability,-runtime,-whitespace,+build/include_subdir,+build/forward_decl,+build/include_what_you_use,+build/deprecated,+whitespace/comma,+whitespace/line_length,+whitespace/empty_if_body,+whitespace/empty_loop_body,+whitespace/empty_conditional_body,+whitespace/forcolon,+whitespace/parens,+whitespace/semicolon,+whitespace/tab,+readability/braces msgq/
 
     # *** tests ***
     test_runner:

msgq/ipc_py.cc

@@ -0,0 +1,117 @@
+#include <cerrno>
+#include <chrono>
+#include <cstdint>
+#include <cstring>
+#include <memory>
+#include <optional>
+#include <stdexcept>
+#include <string>
+#include <vector>
+
+#include <nanobind/nanobind.h>
+#include <nanobind/stl/optional.h>
+#include <nanobind/stl/vector.h>
+
+#include "msgq/event.h"
+#include "msgq/ipc.h"
+#include "msgq/python_helpers.h"
+
+namespace nb = nanobind;
+
+namespace {
+
+struct IpcError : std::runtime_error {
+  using std::runtime_error::runtime_error;
+};
+
+struct MultiplePublishersError : IpcError {
+  using IpcError::IpcError;
+};
+
+void check_error(bool failed, int error, const std::string &endpoint = "") {
+  if (!failed) return;
+  std::string message = "Messaging failure " + (endpoint.empty() ? "" : "with " + endpoint) + ": " + std::strerror(error);
+  if (error == EADDRINUSE) throw MultiplePublishersError(message);
+  throw IpcError(message);
+}
+
+}  // namespace
+
+NB_MODULE(ipc_pyx, m) {
+  nb::exception<IpcError> ipc_error(m, "IpcError");
+  nb::exception<MultiplePublishersError>(m, "MultiplePublishersError", ipc_error);
+
+  m.def("toggle_fake_events", &SocketEventHandle::toggle_fake_events, nb::arg("enabled"));
+  m.def("set_fake_prefix", [](const PythonString &prefix) { SocketEventHandle::set_fake_prefix(python_string(prefix)); }, nb::arg("prefix"));
+  m.def("get_fake_prefix", [] {
+    auto prefix = SocketEventHandle::fake_prefix();
+    return nb::bytes(prefix.data(), prefix.size());
+  });
+  m.def("delete_fake_prefix", [] { SocketEventHandle::set_fake_prefix(""); });
+  m.def("wait_for_one_event", &Event::wait_for_one, nb::arg("events"), nb::arg("timeout") = -1, nb::call_guard<nb::gil_scoped_release>());
+
+  nb::class_<Event>(m, "Event")
+    .def(nb::init<>())
+    .def("set", &Event::set)
+    .def("clear", &Event::clear)
+    .def("wait", &Event::wait, nb::arg("timeout") = -1, nb::call_guard<nb::gil_scoped_release>())
+    .def("peek", &Event::peek)
+    .def_prop_ro("fd", &Event::fd)
+    .def_prop_ro("ptr", [](Event &event) { return reinterpret_cast<uintptr_t>(&event); });
+
+  nb::class_<SocketEventHandle>(m, "SocketEventHandle")
+    .def("__init__", [](SocketEventHandle *self, const PythonString &endpoint, const PythonString &identifier, bool override) {
+      new (self) SocketEventHandle(python_string(endpoint), python_string(identifier), override);
+    }, nb::arg("endpoint"), nb::arg("identifier"), nb::arg("override"))
+    .def_prop_rw("enabled", &SocketEventHandle::is_enabled, &SocketEventHandle::set_enabled)
+    .def_prop_ro("recv_called_event", &SocketEventHandle::recv_called, nb::keep_alive<0, 1>())
+    .def_prop_ro("recv_ready_event", &SocketEventHandle::recv_ready, nb::keep_alive<0, 1>());
+
+  nb::class_<SubSocket>(m, "SubSocket")
+    .def(nb::init<>())
+    .def("connect", [](SubSocket &socket, const PythonString &endpoint, bool conflate, size_t segment_size) {
+      auto name = python_string(endpoint);
+      nb::gil_scoped_release release;
+      int result = socket.connect(name, conflate, segment_size);
+      check_error(result != 0, errno, name);
+    }, nb::arg("endpoint"), nb::arg("conflate") = false, nb::arg("segment_size") = 0)
+    .def("setTimeout", &SubSocket::setTimeout, nb::arg("timeout"), nb::call_guard<nb::gil_scoped_release>())
+    .def("receive", [](SubSocket &socket, bool non_blocking) -> std::optional<nb::bytes> {
+      std::unique_ptr<Message> message;
+      {
+        nb::gil_scoped_release release;
+        message.reset(socket.receive(non_blocking));
+      }
+      if (!message) return std::nullopt;
+      return nb::bytes(message->getData(), message->getSize());
+    }, nb::arg("non_blocking") = false);
+
+  nb::class_<Poller>(m, "Poller")
+    .def(nb::init<>())
+    .def("registerSocket", &Poller::registerSocket, nb::arg("socket"), nb::keep_alive<1, 2>())
+    .def("poll", &Poller::poll, nb::arg("timeout"), nb::rv_policy::reference, nb::call_guard<nb::gil_scoped_release>());
+
+  nb::class_<PubSocket>(m, "PubSocket")
+    .def(nb::init<>())
+    .def("connect", [](PubSocket &socket, const PythonString &endpoint, size_t segment_size) {
+      auto name = python_string(endpoint);
+      int result = socket.connect(name, segment_size);
+      check_error(result != 0, errno, name);
+    }, nb::arg("endpoint"), nb::arg("segment_size") = 0)
+    .def("send", [](PubSocket &socket, nb::bytes data) {
+      int result = socket.send(const_cast<char *>(data.c_str()), data.size());
+      check_error(result != static_cast<int>(data.size()), errno);
+    }, nb::arg("data"))
+    .def("all_readers_updated", &PubSocket::all_readers_updated)
+    .def("wait_for_readers", [](PubSocket &socket, double timeout, double interval) {
+      // Python's sleep releases the GIL and handles signals while waiting.
+      nb::object sleep = nb::module_::import_("time").attr("sleep");
+      auto deadline = std::chrono::steady_clock::now() + std::chrono::duration<double>(timeout);
+      while (std::chrono::steady_clock::now() < deadline) {
+        if (socket.all_readers_updated()) return;
+        sleep(interval);
+      }
+      PyErr_SetString(PyExc_TimeoutError, "subscriber did not connect");
+      throw nb::python_error();
+    }, nb::arg("timeout") = 1.0, nb::arg("interval") = 0.001);
+}