From 78162344748f3afb4b1f746796a893bcec80ea1c Mon Sep 17 00:00:00 2001 From: Adeeb Shihadeh Date: Sun, 25 Oct 2026 12:20:46 -0700 Subject: [PATCH 1/4] Switch to nanobind --- MANIFEST.in | 4 +- README.md | 3 + SConscript | 9 +- SConstruct | 24 ++- lefthook.yml | 6 +- msgq/ipc.pxd | 61 ------- msgq/ipc_py.cc | 119 +++++++++++++ msgq/ipc_pyx.pyi | 77 ++++++++ msgq/ipc_pyx.pyx | 234 ------------------------- msgq/py.typed | 0 msgq/python_helpers.h | 18 ++ msgq/tests/test_bindings.py | 73 ++++++++ msgq/tests/test_typing.py | 43 +++++ msgq/visionipc/tests/test_visionipc.py | 55 ++++++ msgq/visionipc/visionipc.pxd | 52 ------ msgq/visionipc/visionipc_py.cc | 114 ++++++++++++ msgq/visionipc/visionipc_pyx.pxd | 10 -- msgq/visionipc/visionipc_pyx.pyi | 79 +++++++++ msgq/visionipc/visionipc_pyx.pyx | 163 ----------------- pyproject.toml | 15 +- site_scons/site_tools/cython.py | 72 -------- test.sh | 4 +- tools/generate_stubs.py | 31 ++++ 23 files changed, 645 insertions(+), 621 deletions(-) delete mode 100644 msgq/ipc.pxd create mode 100644 msgq/ipc_py.cc create mode 100644 msgq/ipc_pyx.pyi delete mode 100644 msgq/ipc_pyx.pyx create mode 100644 msgq/py.typed create mode 100644 msgq/python_helpers.h create mode 100644 msgq/tests/test_bindings.py create mode 100644 msgq/tests/test_typing.py delete mode 100644 msgq/visionipc/visionipc.pxd create mode 100644 msgq/visionipc/visionipc_py.cc delete mode 100644 msgq/visionipc/visionipc_pyx.pxd create mode 100644 msgq/visionipc/visionipc_pyx.pyi delete mode 100644 msgq/visionipc/visionipc_pyx.pyx delete mode 100644 site_scons/site_tools/cython.py create mode 100644 tools/generate_stubs.py diff --git a/MANIFEST.in b/MANIFEST.in index b3bc16ab9..ae6ac491e 100644 --- a/MANIFEST.in +++ b/MANIFEST.in @@ -1,4 +1,4 @@ include README.md pyproject.toml setup.py SConstruct SConscript -recursive-include site_scons *.py -recursive-include msgq *.py *.pyx *.pxd *.h *.cc +recursive-include tools *.py +recursive-include msgq *.py *.pyi py.typed *.h *.cc global-exclude *.so *.o *.os *.a *.pyc diff --git a/README.md b/README.md index 1609f32a8..12b08aedc 100644 --- a/README.md +++ b/README.md @@ -107,6 +107,9 @@ 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 and ship type stubs for both messaging and vision IPC. After changing a binding, run +`python tools/generate_stubs.py` in the development environment to update the stubs. The test hooks check that they match the bindings and run `ty` against the Python API. + ## License MSGQ is available under the [MIT License](https://github.com/commaai/msgq/blob/master/LICENSE). diff --git a/SConscript b/SConscript index 5a8d52041..a7396d342 100644 --- a/SConscript +++ b/SConscript @@ -1,4 +1,4 @@ -Import('env', 'envCython', 'common') +Import('env', 'envNanobind', 'nanobind_lib', 'common') visionipc_dir = Dir('msgq/visionipc') @@ -10,7 +10,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('msgq/ipc_pyx.so', ['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,8 +25,8 @@ 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 +envNanobind.Program(f'{visionipc_dir.abspath}/visionipc_pyx.so', [f'{visionipc_dir.abspath}/visionipc_py.cc', nanobind_lib], LIBS=vipc_libs) if GetOption('extras'): diff --git a/SConstruct b/SConstruct index b7e1a4c76..d3f383401 100644 --- a/SConstruct +++ b/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']) diff --git a/lefthook.yml b/lefthook.yml index 91b18aa52..708f9d42d 100644 --- a/lefthook.yml +++ b/lefthook.yml @@ -17,13 +17,15 @@ test: run: ruff check . ty: run: ty check . + stubs: + run: python tools/generate_stubs.py --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: diff --git a/msgq/ipc.pxd b/msgq/ipc.pxd deleted file mode 100644 index 3d35af0ae..000000000 --- a/msgq/ipc.pxd +++ /dev/null @@ -1,61 +0,0 @@ -# distutils: language = c++ -#cython: language_level=3 - -from libcpp.string cimport string -from libcpp.vector cimport vector -from libcpp cimport bool - - -cdef extern from "msgq/event.h": - cdef cppclass Event: - @staticmethod - int wait_for_one(vector[Event], int) except + - - Event() - Event(int) - void set() - int clear() - void wait(int) except + - bool peek() - int fd() - - cdef cppclass SocketEventHandle: - @staticmethod - void toggle_fake_events(bool) - @staticmethod - void set_fake_prefix(string) - @staticmethod - string fake_prefix() - - SocketEventHandle(string, string, bool) - bool is_enabled() - void set_enabled(bool) - Event recv_called() - Event recv_ready() - - -cdef extern from "msgq/ipc.h": - cdef cppclass Message: - void init(size_t) - void init(char *, size_t) - void close() - size_t getSize() - char *getData() - - cdef cppclass SubSocket: - SubSocket() nogil - int connect(string, bool, size_t) nogil - Message * receive(bool) nogil - void setTimeout(int) nogil - - cdef cppclass PubSocket: - PubSocket() - int connect(string, size_t) - int sendMessage(Message *) - int send(char *, size_t) - bool all_readers_updated() - - cdef cppclass Poller: - Poller() - void registerSocket(SubSocket *) - vector[SubSocket*] poll(int) nogil diff --git a/msgq/ipc_py.cc b/msgq/ipc_py.cc new file mode 100644 index 000000000..1c6e81a35 --- /dev/null +++ b/msgq/ipc_py.cc @@ -0,0 +1,119 @@ +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include + +#include +#include +#include + +#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 ipc_error(m, "IpcError"); + nb::exception(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::class_(m, "Event") + .def(nb::init<>()) + .def("set", &Event::set) + .def("clear", &Event::clear) + .def("wait", &Event::wait, nb::arg("timeout") = -1, nb::call_guard()) + .def("peek", &Event::peek) + .def_prop_ro("fd", &Event::fd) + .def_prop_ro("ptr", [](Event &event) { return reinterpret_cast(&event); }); + + nb::class_(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_(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()) + .def("receive", [](SubSocket &socket, bool non_blocking) -> std::optional { + std::unique_ptr 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_(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::class_(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(data.c_str()), data.size()); + check_error(result != static_cast(data.size()), errno); + }, nb::arg("data")) + .def("all_readers_updated", &PubSocket::all_readers_updated) + .def("wait_for_readers", [](PubSocket &socket, double timeout, double interval) { + { + nb::gil_scoped_release release; + auto deadline = std::chrono::steady_clock::now() + std::chrono::duration(timeout); + while (std::chrono::steady_clock::now() < deadline) { + if (socket.all_readers_updated()) return; + std::this_thread::sleep_for(std::chrono::duration(interval)); + } + } + PyErr_SetString(PyExc_TimeoutError, "subscriber did not connect"); + throw nb::python_error(); + }, nb::arg("timeout") = 1.0, nb::arg("interval") = 0.001); +} diff --git a/msgq/ipc_pyx.pyi b/msgq/ipc_pyx.pyi new file mode 100644 index 000000000..15e6cecb7 --- /dev/null +++ b/msgq/ipc_pyx.pyi @@ -0,0 +1,77 @@ +from collections.abc import Sequence + + +class IpcError(Exception): + pass + +class MultiplePublishersError(IpcError): + pass + +def toggle_fake_events(enabled: bool) -> None: ... + +def set_fake_prefix(prefix: str | bytes) -> None: ... + +def get_fake_prefix() -> bytes: ... + +def delete_fake_prefix() -> None: ... + +def wait_for_one_event(events: Sequence[Event], timeout: int = -1) -> int: ... + +class Event: + def __init__(self) -> None: ... + + def set(self) -> None: ... + + def clear(self) -> int: ... + + def wait(self, timeout: int = -1) -> None: ... + + def peek(self) -> bool: ... + + @property + def fd(self) -> int: ... + + @property + def ptr(self) -> int: ... + +class SocketEventHandle: + def __init__(self, endpoint: str | bytes, identifier: str | bytes, override: bool) -> None: ... + + @property + def enabled(self) -> bool: ... + + @enabled.setter + def enabled(self, arg: bool, /) -> None: ... + + @property + def recv_called_event(self) -> Event: ... + + @property + def recv_ready_event(self) -> Event: ... + +class SubSocket: + def __init__(self) -> None: ... + + def connect(self, endpoint: str | bytes, conflate: bool = False, segment_size: int = 0) -> None: ... + + def setTimeout(self, timeout: int) -> None: ... + + def receive(self, non_blocking: bool = False) -> bytes | None: ... + +class Poller: + def __init__(self) -> None: ... + + def registerSocket(self, socket: SubSocket) -> None: ... + + def poll(self, timeout: int) -> list[SubSocket]: ... + +class PubSocket: + def __init__(self) -> None: ... + + def connect(self, endpoint: str | bytes, segment_size: int = 0) -> None: ... + + def send(self, data: bytes) -> None: ... + + def all_readers_updated(self) -> bool: ... + + def wait_for_readers(self, timeout: float = 1.0, interval: float = 0.001) -> None: ... diff --git a/msgq/ipc_pyx.pyx b/msgq/ipc_pyx.pyx deleted file mode 100644 index 446f303f6..000000000 --- a/msgq/ipc_pyx.pyx +++ /dev/null @@ -1,234 +0,0 @@ -# distutils: language = c++ -# cython: c_string_encoding=ascii, language_level=3 - -import time -from libcpp.string cimport string -from libcpp.vector cimport vector -from libcpp cimport bool -from libc cimport errno -from libc.string cimport strerror -from cython.operator import dereference - - -from .ipc cimport SubSocket as cppSubSocket -from .ipc cimport PubSocket as cppPubSocket -from .ipc cimport Poller as cppPoller -from .ipc cimport Message as cppMessage -from .ipc cimport Event as cppEvent, SocketEventHandle as cppSocketEventHandle - - -class IpcError(Exception): - def __init__(self, endpoint=None): - suffix = f"with {endpoint.decode('utf-8')}" if endpoint else "" - message = f"Messaging failure {suffix}: {strerror(errno.errno).decode('utf-8')}" - super().__init__(message) - - -class MultiplePublishersError(IpcError): - pass - - -def toggle_fake_events(bool enabled): - cppSocketEventHandle.toggle_fake_events(enabled) - - -def set_fake_prefix(string prefix): - cppSocketEventHandle.set_fake_prefix(prefix) - - -def get_fake_prefix(): - return cppSocketEventHandle.fake_prefix() - - -def delete_fake_prefix(): - cppSocketEventHandle.set_fake_prefix(b"") - - -def wait_for_one_event(list events, int timeout=-1): - cdef vector[cppEvent] items - for event in events: - items.push_back(dereference(event.ptr)) - return cppEvent.wait_for_one(items, timeout) - - -cdef class Event: - cdef cppEvent event; - - cdef setEvent(self, cppEvent event): - self.event = event - - def set(self): - self.event.set() - - def clear(self): - return self.event.clear() - - def wait(self, int timeout=-1): - self.event.wait(timeout) - - def peek(self): - return self.event.peek() - - @property - def fd(self): - return self.event.fd() - - @property - def ptr(self): - return &self.event - - -cdef class SocketEventHandle: - cdef cppSocketEventHandle * handle; - - def __cinit__(self, string endpoint, string identifier, bool override): - self.handle = new cppSocketEventHandle(endpoint, identifier, override) - - def __dealloc__(self): - del self.handle - - @property - def enabled(self): - return self.handle.is_enabled() - - @enabled.setter - def enabled(self, bool value): - self.handle.set_enabled(value) - - @property - def recv_called_event(self): - e = Event() - e.setEvent(self.handle.recv_called()) - - return e - - @property - def recv_ready_event(self): - e = Event() - e.setEvent(self.handle.recv_ready()) - - return e - - -cdef class Poller: - cdef cppPoller * poller - cdef list sub_sockets - - def __cinit__(self): - self.sub_sockets = [] - self.poller = new cppPoller() - - def __dealloc__(self): - del self.poller - - def registerSocket(self, SubSocket socket): - self.sub_sockets.append(socket) - self.poller.registerSocket(socket.socket) - - def poll(self, timeout): - sockets = [] - cdef int t = timeout - - with nogil: - result = self.poller.poll(t) - - for s in result: - socket = SubSocket() - socket.setPtr(s) - sockets.append(socket) - - return sockets - - -cdef class SubSocket: - cdef cppSubSocket * socket - cdef bool is_owner - - def __cinit__(self): - with nogil: - self.socket = new cppSubSocket() - - self.is_owner = True - - def __dealloc__(self): - if self.is_owner: - with nogil: - del self.socket - - cdef setPtr(self, cppSubSocket * ptr): - if self.is_owner: - with nogil: - del self.socket - - self.is_owner = False - self.socket = ptr - - def connect(self, string endpoint, bool conflate=False, size_t segment_size=0): - cdef int r - with nogil: - r = self.socket.connect(endpoint, conflate, segment_size) - - if r != 0: - if errno.errno == errno.EADDRINUSE: - raise MultiplePublishersError(endpoint) - else: - raise IpcError(endpoint) - - def setTimeout(self, int timeout): - with nogil: - self.socket.setTimeout(timeout) - - def receive(self, bool non_blocking=False): - cdef cppMessage *msg - with nogil: - msg = self.socket.receive(non_blocking) - - if msg == NULL: - return None - else: - sz = msg.getSize() - m = msg.getData()[:sz] - with nogil: - del msg - - return m - - -cdef class PubSocket: - cdef cppPubSocket * socket - - def __cinit__(self): - self.socket = new cppPubSocket() - - def __dealloc__(self): - del self.socket - - def connect(self, string endpoint, size_t segment_size=0): - r = self.socket.connect(endpoint, segment_size) - - if r != 0: - if errno.errno == errno.EADDRINUSE: - raise MultiplePublishersError(endpoint) - else: - raise IpcError(endpoint) - - def send(self, bytes data): - length = len(data) - r = self.socket.send(data, length) - - if r != length: - if errno.errno == errno.EADDRINUSE: - raise MultiplePublishersError - else: - raise IpcError - - def all_readers_updated(self): - return self.socket.all_readers_updated() - - def wait_for_readers(self, double timeout=1.0, double interval=0.001): - deadline = time.monotonic() + timeout - while time.monotonic() < deadline: - if self.all_readers_updated(): - return - time.sleep(interval) - raise TimeoutError("subscriber did not connect") diff --git a/msgq/py.typed b/msgq/py.typed new file mode 100644 index 000000000..e69de29bb diff --git a/msgq/python_helpers.h b/msgq/python_helpers.h new file mode 100644 index 000000000..291230111 --- /dev/null +++ b/msgq/python_helpers.h @@ -0,0 +1,18 @@ +#pragma once + +#include +#include + +#include +#include +#include + +// Cython accepted both text and bytes for C++ strings. Keep that API while +// exposing the full input type to nanobind's stub generator. +using PythonString = std::variant; + +inline std::string python_string(const PythonString &value) { + if (const auto *text = std::get_if(&value)) return *text; + const auto &bytes = std::get(value); + return std::string(bytes.c_str(), bytes.size()); +} diff --git a/msgq/tests/test_bindings.py b/msgq/tests/test_bindings.py new file mode 100644 index 000000000..c7f39c786 --- /dev/null +++ b/msgq/tests/test_bindings.py @@ -0,0 +1,73 @@ +import gc +import threading +import unittest + +import msgq +from msgq.ipc_pyx import Event +from msgq.tests.test_messaging import random_sock + + +class TestBindings(unittest.TestCase): + def test_bytes_endpoints_and_binary_messages(self): + endpoint = random_sock().encode() + publisher = msgq.PubSocket() + subscriber = msgq.SubSocket() + publisher.connect(endpoint) + subscriber.connect(endpoint) + publisher.send(b"\x00\xffpayload\x00") + self.assertEqual(subscriber.receive(), b"\x00\xffpayload\x00") + with self.assertRaises(TypeError): + getattr(publisher, "send")("text") + + def test_fake_prefix_bytes(self): + try: + for prefix in ("prefix", b"prefix"): + msgq.set_fake_prefix(prefix) + self.assertEqual(msgq.get_fake_prefix(), prefix.encode() if isinstance(prefix, str) else prefix) + finally: + msgq.delete_fake_prefix() + + def test_event_keeps_handle_alive(self): + event = msgq.fake_event_handle(random_sock()).recv_called_event + gc.collect() + event.set() + self.assertTrue(event.peek()) + event.clear() + self.assertFalse(event.peek()) + + def test_event_wait_releases_gil(self): + handle = msgq.fake_event_handle(random_sock()) + event = handle.recv_called_event + for wait in (lambda: event.wait(1), lambda: msgq.wait_for_one_event([event], 1)): + with self.subTest(wait=wait): + event.clear() + timer = threading.Timer(0.01, event.set) + timer.start() + try: + wait() + self.assertTrue(event.peek()) + finally: + timer.join() + + def test_poller_socket_lifetime(self): + endpoint = random_sock() + publisher = msgq.pub_sock(endpoint) + poller = msgq.Poller() + subscriber = msgq.sub_sock(endpoint, poller=poller) + del subscriber + gc.collect() + publisher.send(b"message") + sockets = poller.poll(1000) + self.assertEqual(len(sockets), 1) + del poller + gc.collect() + self.assertEqual(sockets[0].receive(), b"message") + + def test_reader_timeout(self): + publisher = msgq.pub_sock(random_sock()) + with self.assertRaisesRegex(TimeoutError, "subscriber did not connect"): + publisher.wait_for_readers(timeout=0) + + def test_invalid_event(self): + with self.assertRaisesRegex(RuntimeError, "valid file descriptor"): + Event().wait(0) diff --git a/msgq/tests/test_typing.py b/msgq/tests/test_typing.py new file mode 100644 index 000000000..4f48b011d --- /dev/null +++ b/msgq/tests/test_typing.py @@ -0,0 +1,43 @@ +from typing import TYPE_CHECKING, assert_type + +if TYPE_CHECKING: + import msgq + from msgq.ipc_pyx import Event + from msgq.visionipc import VisionBuf, VisionIpcClient, VisionIpcServer, get_endpoint_name + + publisher = msgq.pub_sock("typing") + subscriber = msgq.sub_sock("typing") + poller = msgq.Poller() + poller.registerSocket(subscriber) + assert_type(publisher.send(b"data"), None) + assert_type(publisher.all_readers_updated(), bool) + assert_type(subscriber.receive(), bytes | None) + assert_type(poller.poll(0), list[msgq.SubSocket]) + assert_type(msgq.drain_sock_raw(subscriber), list[bytes]) + assert_type(msgq.get_fake_prefix(), bytes) + + handle = msgq.fake_event_handle("typing", identifier=b"prefix") + handle.enabled = True + assert_type(handle.enabled, bool) + assert_type(handle.recv_called_event, Event) + assert_type(handle.recv_ready_event.fd, int) + assert_type(msgq.wait_for_one_event([handle.recv_called_event], 0), int) + + server = VisionIpcServer("typing") + client = VisionIpcClient("typing", 0, False) + assert_type(get_endpoint_name("typing", 0), str) + assert_type(VisionIpcClient.available_streams("typing", False), set[int]) + assert_type(client.connect(False), bool) + assert_type(client.width, int | None) + assert_type(client.buffer_len, int | None) + assert_type(client.recv(), VisionBuf | None) + assert_type(client.frame_id, int) + assert_type(client.valid, bool) + frame = client.recv() + if frame is not None: + assert_type(frame.data, memoryview) + assert_type(frame.width, int) + assert_type(frame.frame_id, int) + server.send(0, frame.data) + server.send(0, b"frame") + server.send(0, bytearray(b"frame")) diff --git a/msgq/visionipc/tests/test_visionipc.py b/msgq/visionipc/tests/test_visionipc.py index b938df447..9239e9d66 100644 --- a/msgq/visionipc/tests/test_visionipc.py +++ b/msgq/visionipc/tests/test_visionipc.py @@ -1,3 +1,4 @@ +import gc import struct import unittest from typing import Optional @@ -101,3 +102,57 @@ def test_conflate(self): recv_buf = self.client.recv(timeout_ms=5) assert recv_buf is None + + def test_send_buffer_types(self): + self.setup_vipc("camerad", 0) + assert self.server is not None + assert self.client is not None + assert self.client.buffer_len is not None + data = bytes(i % 256 for i in range(self.client.buffer_len)) + strided = bytearray(len(data) * 2) + strided[::2] = data + for payload in (data, bytearray(data), memoryview(data), memoryview(strided)[::2], memoryview(data[::-1])[::-1]): + with self.subTest(type=type(payload)): + self.server.send(0, payload, frame_id=42, timestamp_sof=123, timestamp_eof=456) + buf = self.client.recv() + assert buf is not None + self.assertEqual(bytes(buf.data), data) + self.assertEqual(buf.frame_id, 42) + self.assertEqual(self.client.timestamp_sof, 123) + self.assertEqual(self.client.timestamp_eof, 456) + self.assertFalse(self.client.valid) + + def test_send_wrong_size(self): + self.setup_vipc("camerad", 0) + assert self.server is not None + with self.assertRaises(AssertionError): + self.server.send(0, b"short") + + def test_buffer_keeps_client_alive(self): + self.setup_vipc("camerad", 0) + assert self.server is not None + assert self.client is not None + assert self.client.buffer_len is not None + data = b"a" * self.client.buffer_len + self.server.send(0, data, frame_id=42) + buf = self.client.recv() + assert buf is not None + self.client = None + gc.collect() + self.assertEqual(buf.frame_id, 42) + view = buf.data + del buf + gc.collect() + self.assertEqual(bytes(view), data) + self.assertFalse(view.readonly) + view[0] = ord("b") + self.assertEqual(view[0], ord("b")) + + def test_unconnected_properties(self): + client = VisionIpcClient("unconnected", 0, False) + for value in (client.width, client.height, client.stride, client.uv_offset, client.buffer_len): + self.assertIsNone(value) + self.assertEqual(client.num_buffers, 0) + self.assertFalse(client.is_connected()) + self.assertFalse(client.connect(False)) + self.assertIsNone(client.recv(timeout_ms=0)) diff --git a/msgq/visionipc/visionipc.pxd b/msgq/visionipc/visionipc.pxd deleted file mode 100644 index 2552fe0e4..000000000 --- a/msgq/visionipc/visionipc.pxd +++ /dev/null @@ -1,52 +0,0 @@ -# distutils: language = c++ -#cython: language_level=3 - -from libcpp.string cimport string -from libcpp.vector cimport vector -from libcpp.set cimport set -from libc.stdint cimport uint32_t, uint64_t -from libcpp cimport bool, int - -cdef extern from "msgq/visionipc/visionbuf.h": - ctypedef uint32_t VisionStreamType "VisionStreamType" - - cdef cppclass VisionBuf: - void * addr - int fd - size_t len - size_t width - size_t height - size_t stride - size_t uv_offset - size_t idx - void set_frame_id(uint64_t id) - uint64_t get_frame_id() - -cdef extern from "msgq/visionipc/visionipc.h": - struct VisionIpcBufExtra: - uint32_t frame_id - uint64_t timestamp_sof - uint64_t timestamp_eof - bool valid - -cdef extern from "msgq/visionipc/visionipc_server.h": - string get_endpoint_name(string, VisionStreamType) - - cdef cppclass VisionIpcServer: - VisionIpcServer(string) - void create_buffers(VisionStreamType, size_t, size_t, size_t) - void create_buffers_with_sizes(VisionStreamType, size_t, size_t, size_t, size_t, size_t, size_t) - VisionBuf * get_buffer(VisionStreamType) - void send(VisionBuf *, VisionIpcBufExtra *, bool) - void start_listener() - -cdef extern from "msgq/visionipc/visionipc_client.h": - cdef cppclass VisionIpcClient: - int num_buffers - VisionBuf buffers[1] - VisionIpcClient(string, VisionStreamType, bool) - VisionBuf * recv(VisionIpcBufExtra *, int) nogil - bool connect(bool) - bool is_connected() - @staticmethod - set[VisionStreamType] getAvailableStreams(string, bool) diff --git a/msgq/visionipc/visionipc_py.cc b/msgq/visionipc/visionipc_py.cc new file mode 100644 index 000000000..c6b7a6c65 --- /dev/null +++ b/msgq/visionipc/visionipc_py.cc @@ -0,0 +1,114 @@ +#include +#include +#include +#include + +#include +#include +#include +#include + +#include "msgq/python_helpers.h" +#include "msgq/visionipc/visionbuf.h" +#include "msgq/visionipc/visionipc.h" +#include "msgq/visionipc/visionipc_client.h" +#include "msgq/visionipc/visionipc_server.h" + +namespace nb = nanobind; + +namespace { + +struct PythonVisionIpcClient : VisionIpcClient { + using VisionIpcClient::VisionIpcClient; + VisionIpcBufExtra extra = {}; +}; + +using Buffer = nb::ndarray, nb::device::cpu>; +using InputBuffer = nb::ndarray, nb::device::cpu>; + +} // namespace + +NB_MODULE(visionipc_pyx, m) { + m.def("get_endpoint_name", [](const PythonString &name, VisionStreamType stream) { + return get_endpoint_name(python_string(name), stream); + }, nb::arg("name"), nb::arg("stream")); + + nb::class_(m, "VisionBuf") + .def_prop_ro("data", [](VisionBuf &buf) { return Buffer(buf.addr, {buf.len}); }, nb::rv_policy::reference_internal, + nb::sig("def data(self) -> memoryview")) + .def_ro("width", &VisionBuf::width) + .def_ro("height", &VisionBuf::height) + .def_ro("stride", &VisionBuf::stride) + .def_ro("uv_offset", &VisionBuf::uv_offset) + .def_ro("idx", &VisionBuf::idx) + .def_ro("fd", &VisionBuf::fd) + .def_prop_ro("frame_id", &VisionBuf::get_frame_id); + + nb::class_(m, "VisionIpcServer") + .def("__init__", [](VisionIpcServer *self, const PythonString &name) { + new (self) VisionIpcServer(python_string(name)); + }, nb::arg("name")) + .def("create_buffers", &VisionIpcServer::create_buffers, nb::arg("tp"), nb::arg("num_buffers"), nb::arg("width"), nb::arg("height")) + .def("create_buffers_with_sizes", &VisionIpcServer::create_buffers_with_sizes, + nb::arg("tp"), nb::arg("num_buffers"), nb::arg("width"), nb::arg("height"), nb::arg("size"), nb::arg("stride"), nb::arg("uv_offset")) + .def("send", [](VisionIpcServer &server, VisionStreamType tp, InputBuffer data, + uint32_t frame_id, uint64_t timestamp_sof, uint64_t timestamp_eof) { + VisionBuf *buf = server.get_buffer(tp); + if (buf->len != data.size()) { + PyErr_SetString(PyExc_AssertionError, "buffer size does not match frame size"); + throw nb::python_error(); + } + // The old binding accepted strided byte views as well as contiguous data. + auto *source = data.data(); + if (data.stride(0) == 1) { + std::memcpy(buf->addr, source, data.size()); + } else { + auto *target = static_cast(buf->addr); + for (size_t i = 0; i < data.size(); ++i) target[i] = source[static_cast(i) * data.stride(0)]; + } + buf->set_frame_id(frame_id); + VisionIpcBufExtra extra = {frame_id, timestamp_sof, timestamp_eof, false}; + server.send(buf, &extra, false); + }, nb::arg("tp"), nb::arg("data").noconvert(), nb::arg("frame_id") = 0, nb::arg("timestamp_sof") = 0, nb::arg("timestamp_eof") = 0, + nb::sig("def send(self, tp: int, data: bytes | bytearray | memoryview, frame_id: int = 0, timestamp_sof: int = 0, timestamp_eof: int = 0) -> None")) + .def("start_listener", &VisionIpcServer::start_listener); + + nb::class_(m, "VisionIpcClient") + .def("__init__", [](PythonVisionIpcClient *self, const PythonString &name, VisionStreamType stream, bool conflate) { + new (self) PythonVisionIpcClient(python_string(name), stream, conflate); + }, nb::arg("name"), nb::arg("stream"), nb::arg("conflate")) + .def_prop_ro("width", [](PythonVisionIpcClient &client) -> std::optional { + return client.num_buffers ? std::make_optional(client.buffers[0].width) : std::nullopt; + }) + .def_prop_ro("height", [](PythonVisionIpcClient &client) -> std::optional { + return client.num_buffers ? std::make_optional(client.buffers[0].height) : std::nullopt; + }) + .def_prop_ro("stride", [](PythonVisionIpcClient &client) -> std::optional { + return client.num_buffers ? std::make_optional(client.buffers[0].stride) : std::nullopt; + }) + .def_prop_ro("uv_offset", [](PythonVisionIpcClient &client) -> std::optional { + return client.num_buffers ? std::make_optional(client.buffers[0].uv_offset) : std::nullopt; + }) + .def_prop_ro("buffer_len", [](PythonVisionIpcClient &client) -> std::optional { + return client.num_buffers ? std::make_optional(client.buffers[0].len) : std::nullopt; + }) + .def_prop_ro("num_buffers", [](PythonVisionIpcClient &client) { return client.num_buffers; }) + .def_prop_ro("frame_id", [](PythonVisionIpcClient &client) { return client.extra.frame_id; }) + .def_prop_ro("timestamp_sof", [](PythonVisionIpcClient &client) { return client.extra.timestamp_sof; }) + .def_prop_ro("timestamp_eof", [](PythonVisionIpcClient &client) { return client.extra.timestamp_eof; }) + .def_prop_ro("valid", [](PythonVisionIpcClient &client) { return client.extra.valid; }) + .def("recv", [](PythonVisionIpcClient &client, int timeout_ms) { + nb::gil_scoped_release release; + return client.recv(&client.extra, timeout_ms); + }, nb::arg("timeout_ms") = 100, nb::rv_policy::reference_internal, + nb::sig("def recv(self, timeout_ms: int = 100) -> VisionBuf | None")) + .def("connect", [](PythonVisionIpcClient &client, bool blocking) { + return client.connect(blocking); + }, nb::arg("blocking"), nb::call_guard()) + .def("is_connected", [](PythonVisionIpcClient &client) { return client.is_connected(); }) + .def_static("available_streams", [](const PythonString &name, bool block) { + auto text = python_string(name); + nb::gil_scoped_release release; + return VisionIpcClient::getAvailableStreams(text, block); + }, nb::arg("name"), nb::arg("block")); +} diff --git a/msgq/visionipc/visionipc_pyx.pxd b/msgq/visionipc/visionipc_pyx.pxd deleted file mode 100644 index 6e2d5ed05..000000000 --- a/msgq/visionipc/visionipc_pyx.pxd +++ /dev/null @@ -1,10 +0,0 @@ -# distutils: language = c++ -#cython: language_level=3 - -from .visionipc cimport VisionBuf as cppVisionBuf - -cdef class VisionBuf: - cdef cppVisionBuf * buf - - @staticmethod - cdef create(cppVisionBuf*) diff --git a/msgq/visionipc/visionipc_pyx.pyi b/msgq/visionipc/visionipc_pyx.pyi new file mode 100644 index 000000000..767383fe2 --- /dev/null +++ b/msgq/visionipc/visionipc_pyx.pyi @@ -0,0 +1,79 @@ +def get_endpoint_name(name: str | bytes, stream: int) -> str: ... + +class VisionBuf: + @property + def data(self) -> memoryview: ... + + @property + def width(self) -> int: ... + + @property + def height(self) -> int: ... + + @property + def stride(self) -> int: ... + + @property + def uv_offset(self) -> int: ... + + @property + def idx(self) -> int: ... + + @property + def fd(self) -> int: ... + + @property + def frame_id(self) -> int: ... + +class VisionIpcServer: + def __init__(self, name: str | bytes) -> None: ... + + def create_buffers(self, tp: int, num_buffers: int, width: int, height: int) -> None: ... + + def create_buffers_with_sizes(self, tp: int, num_buffers: int, width: int, height: int, size: int, stride: int, uv_offset: int) -> None: ... + + def send(self, tp: int, data: bytes | bytearray | memoryview, frame_id: int = 0, timestamp_sof: int = 0, timestamp_eof: int = 0) -> None: ... + + def start_listener(self) -> None: ... + +class VisionIpcClient: + def __init__(self, name: str | bytes, stream: int, conflate: bool) -> None: ... + + @property + def width(self) -> int | None: ... + + @property + def height(self) -> int | None: ... + + @property + def stride(self) -> int | None: ... + + @property + def uv_offset(self) -> int | None: ... + + @property + def buffer_len(self) -> int | None: ... + + @property + def num_buffers(self) -> int: ... + + @property + def frame_id(self) -> int: ... + + @property + def timestamp_sof(self) -> int: ... + + @property + def timestamp_eof(self) -> int: ... + + @property + def valid(self) -> bool: ... + + def recv(self, timeout_ms: int = 100) -> VisionBuf | None: ... + + def connect(self, blocking: bool) -> bool: ... + + def is_connected(self) -> bool: ... + + @staticmethod + def available_streams(name: str | bytes, block: bool) -> set[int]: ... diff --git a/msgq/visionipc/visionipc_pyx.pyx b/msgq/visionipc/visionipc_pyx.pyx deleted file mode 100644 index a274d4f7a..000000000 --- a/msgq/visionipc/visionipc_pyx.pyx +++ /dev/null @@ -1,163 +0,0 @@ -# distutils: language = c++ -# cython: c_string_encoding=ascii, language_level=3 - -from libc.string cimport memcpy -from libc.stdint cimport uint32_t, uint64_t -from libcpp cimport bool -from libcpp.string cimport string - -from .visionipc cimport VisionIpcServer as cppVisionIpcServer -from .visionipc cimport VisionIpcClient as cppVisionIpcClient -from .visionipc cimport VisionBuf as cppVisionBuf -from .visionipc cimport VisionIpcBufExtra, VisionStreamType -from .visionipc cimport get_endpoint_name as cpp_get_endpoint_name - - -def get_endpoint_name(string name, VisionStreamType stream): - return cpp_get_endpoint_name(name, stream).decode('utf-8') - - -cdef class VisionBuf: - @staticmethod - cdef create(cppVisionBuf * cbuf): - buf = VisionBuf() - buf.buf = cbuf - return buf - - @property - def data(self): - cdef unsigned char[:] data = self.buf.addr - return memoryview(data) - - @property - def width(self): - return self.buf.width - - @property - def height(self): - return self.buf.height - - @property - def stride(self): - return self.buf.stride - - @property - def uv_offset(self): - return self.buf.uv_offset - - @property - def idx(self): - return self.buf.idx - - @property - def fd(self): - return self.buf.fd - - @property - def frame_id(self): - return self.buf.get_frame_id() - - -cdef class VisionIpcServer: - cdef cppVisionIpcServer * server - - def __init__(self, string name): - self.server = new cppVisionIpcServer(name) - - def create_buffers(self, VisionStreamType tp, size_t num_buffers, size_t width, size_t height): - self.server.create_buffers(tp, num_buffers, width, height) - - def create_buffers_with_sizes(self, VisionStreamType tp, size_t num_buffers, size_t width, size_t height, size_t size, size_t stride, size_t uv_offset): - self.server.create_buffers_with_sizes(tp, num_buffers, width, height, size, stride, uv_offset) - - def send(self, VisionStreamType tp, const unsigned char[:] data, uint32_t frame_id=0, uint64_t timestamp_sof=0, uint64_t timestamp_eof=0): - cdef cppVisionBuf * buf = self.server.get_buffer(tp) - - # Populate buffer - assert buf.len == len(data) - memcpy(buf.addr, &data[0], len(data)) - buf.set_frame_id(frame_id) - - cdef VisionIpcBufExtra extra - extra.frame_id = frame_id - extra.timestamp_sof = timestamp_sof - extra.timestamp_eof = timestamp_eof - extra.valid = False - - self.server.send(buf, &extra, False) - - def start_listener(self): - self.server.start_listener() - - def __dealloc__(self): - del self.server - - -cdef class VisionIpcClient: - cdef cppVisionIpcClient * client - cdef VisionIpcBufExtra extra - - def __cinit__(self, string name, VisionStreamType stream, bool conflate): - self.client = new cppVisionIpcClient(name, stream, conflate) - - def __dealloc__(self): - del self.client - - @property - def width(self): - return self.client.buffers[0].width if self.client.num_buffers else None - - @property - def height(self): - return self.client.buffers[0].height if self.client.num_buffers else None - - @property - def stride(self): - return self.client.buffers[0].stride if self.client.num_buffers else None - - @property - def uv_offset(self): - return self.client.buffers[0].uv_offset if self.client.num_buffers else None - - @property - def buffer_len(self): - return self.client.buffers[0].len if self.client.num_buffers else None - - @property - def num_buffers(self): - return self.client.num_buffers - - @property - def frame_id(self): - return self.extra.frame_id - - @property - def timestamp_sof(self): - return self.extra.timestamp_sof - - @property - def timestamp_eof(self): - return self.extra.timestamp_eof - - @property - def valid(self): - return self.extra.valid - - def recv(self, int timeout_ms=100): - cdef cppVisionBuf * buf - # release the GIL, this can block for timeout_ms - with nogil: - buf = self.client.recv(&self.extra, timeout_ms) - if not buf: - return None - return VisionBuf.create(buf) - - def connect(self, bool blocking): - return self.client.connect(blocking) - - def is_connected(self): - return self.client.is_connected() - - @staticmethod - def available_streams(string name, bool block): - return cppVisionIpcClient.getAvailableStreams(name, block) diff --git a/pyproject.toml b/pyproject.toml index e2f4d16ce..2735b6fcb 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -1,5 +1,5 @@ [build-system] -requires = ["setuptools>=77", "Cython>=3.1", "scons"] +requires = ["setuptools>=77", "nanobind>=2.12", "scons"] build-backend = "setuptools.build_meta" [project] @@ -23,7 +23,7 @@ Repository = "https://github.com/commaai/msgq" [project.optional-dependencies] dev = [ "setuptools", # for distutils - "Cython", + "nanobind>=2.12", "scons", "ruff", "cppcheck", @@ -41,8 +41,8 @@ include = ["msgq", "msgq.*"] namespaces = false [tool.setuptools.package-data] -"msgq" = ["*.pxd", "*.h"] -"msgq.visionipc" = ["*.pxd", "*.h"] +"msgq" = ["*.pyi", "py.typed", "*.h"] +"msgq.visionipc" = ["*.pyi", "*.h"] [tool.cibuildwheel] build = "cp311-* cp312-* cp313-* cp314-*" @@ -56,10 +56,3 @@ lint.flake8-implicit-str-concat.allow-multiline=false line-length = 160 target-version="py311" - -[tool.ty.src] -exclude = ["site_scons/"] - -[tool.ty.rules] -# Cython modules are compiled at build time, not available for static analysis -unresolved-import = "ignore" diff --git a/site_scons/site_tools/cython.py b/site_scons/site_tools/cython.py deleted file mode 100644 index c29147553..000000000 --- a/site_scons/site_tools/cython.py +++ /dev/null @@ -1,72 +0,0 @@ -import re -import SCons -from SCons.Action import Action -from SCons.Scanner import Scanner - -pyx_from_import_re = re.compile(r'^from\s+(\S+)\s+cimport', re.M) -pyx_import_re = re.compile(r'^cimport\s+(\S+)', re.M) -cdef_import_re = re.compile(r'^cdef extern from\s+.(\S+).:', re.M) - - -def pyx_scan(node, env, path, arg=None): - contents = node.get_text_contents() - - # from cimport ... - matches = pyx_from_import_re.findall(contents) - # cimport - matches += pyx_import_re.findall(contents) - - # Modules can be either .pxd or .pyx files - files = [m.replace('.', '/') + '.pxd' for m in matches] - files += [m.replace('.', '/') + '.pyx' for m in matches] - - # cdef extern from - files += cdef_import_re.findall(contents) - - # Handle relative imports - cur_dir = str(node.get_dir()) - files = [cur_dir + f if f.startswith('/') else f for f in files] - - # Filter out non-existing files (probably system imports) - files = [f for f in files if env.File(f).exists()] - return env.File(files) - - -pyxscanner = Scanner(function=pyx_scan, skeys=['.pyx', '.pxd'], recursive=True) -cythonAction = Action("$CYTHONCOM") - - -def create_builder(env): - try: - cython = env['BUILDERS']['Cython'] - except KeyError: - cython = SCons.Builder.Builder( - action=cythonAction, - emitter={}, - suffix=cython_suffix_emitter, - single_source=1 - ) - env.Append(SCANNERS=pyxscanner) - env['BUILDERS']['Cython'] = cython - return cython - -def cython_suffix_emitter(env, source): - return "$CYTHONCFILESUFFIX" - -def generate(env): - env["CYTHON"] = "cythonize" - env["CYTHONCOM"] = "$CYTHON $CYTHONFLAGS $SOURCE" - env["CYTHONCFILESUFFIX"] = ".cpp" - - c_file, _ = SCons.Tool.createCFileBuilders(env) - - c_file.suffix['.pyx'] = cython_suffix_emitter - c_file.add_action('.pyx', cythonAction) - - c_file.suffix['.py'] = cython_suffix_emitter - c_file.add_action('.py', cythonAction) - - create_builder(env) - -def exists(env): - return True diff --git a/test.sh b/test.sh index bf3a0e0ef..2feab797a 100755 --- a/test.sh +++ b/test.sh @@ -19,7 +19,9 @@ lefthook run test uv venv --python "$DIR/.venv/bin/python" "$TEST_DIR/.venv" uv pip install --python "$TEST_DIR/.venv/bin/python" "$DIR" cd "$TEST_DIR" - "$TEST_DIR/.venv/bin/python" -m unittest msgq.tests.test_messaging + "$TEST_DIR/.venv/bin/python" -m unittest msgq.tests.test_messaging msgq.tests.test_bindings msgq.visionipc.tests.test_visionipc + cp "$DIR/msgq/tests/test_typing.py" "$TEST_DIR/check_types.py" + "$DIR/.venv/bin/ty" check --python "$TEST_DIR/.venv/bin/python" "$TEST_DIR/check_types.py" ) # *** all done *** diff --git a/tools/generate_stubs.py b/tools/generate_stubs.py new file mode 100644 index 000000000..5d0f75867 --- /dev/null +++ b/tools/generate_stubs.py @@ -0,0 +1,31 @@ +import argparse +from pathlib import Path +import subprocess +import sys +import tempfile + +ROOT = Path(__file__).resolve().parents[1] +MODULES = ("msgq.ipc_pyx", "msgq.visionipc.visionipc_pyx") + + +def main() -> None: + parser = argparse.ArgumentParser(description="Generate nanobind stubs, or check that committed stubs are current.") + parser.add_argument("--check", action="store_true") + args = parser.parse_args() + + with tempfile.TemporaryDirectory() as directory: + for module in MODULES: + target = ROOT.joinpath(*module.split(".")).with_suffix(".pyi") + output = Path(directory) / target.name if args.check else target + subprocess.check_call([sys.executable, "-m", "nanobind.stubgen", "-i", str(ROOT), "-m", module, "-o", str(output), "-q"]) + # Older nanobind versions leave blank lines when a stub has no imports. + content = output.read_text().lstrip() + if args.check: + if content != target.read_text(): + sys.exit(f"{target.relative_to(ROOT)} is out of date; run python tools/generate_stubs.py") + else: + output.write_text(content) + + +if __name__ == "__main__": + main() From 549a7a2658ec1e03e573be590396ce878075a0f3 Mon Sep 17 00:00:00 2001 From: Adeeb Shihadeh Date: Sun, 25 Oct 2026 12:24:20 -0700 Subject: [PATCH 2/4] Generate type stubs during build --- .gitignore | 3 ++ MANIFEST.in | 5 +- README.md | 3 +- SConscript | 9 +++- lefthook.yml | 2 - msgq/ipc_pyx.pyi | 77 ------------------------------- msgq/tests/test_bindings.py | 73 ----------------------------- msgq/tests/test_typing.py | 43 ----------------- msgq/visionipc/visionipc_pyx.pyi | 79 -------------------------------- setup.py | 2 + test.sh | 4 +- tools/generate_stubs.py | 31 ------------- 12 files changed, 17 insertions(+), 314 deletions(-) delete mode 100644 msgq/ipc_pyx.pyi delete mode 100644 msgq/tests/test_bindings.py delete mode 100644 msgq/tests/test_typing.py delete mode 100644 msgq/visionipc/visionipc_pyx.pyi delete mode 100644 tools/generate_stubs.py diff --git a/.gitignore b/.gitignore index 9294b1bb5..68de354da 100644 --- a/.gitignore +++ b/.gitignore @@ -20,3 +20,6 @@ libmessaging.* libmessaging_shared.* .sconsign.dblite .mypy_cache/ + +/msgq/ipc_pyx.pyi +/msgq/visionipc/visionipc_pyx.pyi diff --git a/MANIFEST.in b/MANIFEST.in index ae6ac491e..5f694e2a0 100644 --- a/MANIFEST.in +++ b/MANIFEST.in @@ -1,4 +1,3 @@ include README.md pyproject.toml setup.py SConstruct SConscript -recursive-include tools *.py -recursive-include msgq *.py *.pyi py.typed *.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 diff --git a/README.md b/README.md index 12b08aedc..015caf7a8 100644 --- a/README.md +++ b/README.md @@ -107,8 +107,7 @@ 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 and ship type stubs for both messaging and vision IPC. After changing a binding, run -`python tools/generate_stubs.py` in the development environment to update the stubs. The test hooks check that they match the bindings and run `ty` against the Python API. +The Python bindings use nanobind. SCons generates type stubs from the compiled modules for type checking and includes them in the Python package. ## License diff --git a/SConscript b/SConscript index a7396d342..141d2c972 100644 --- a/SConscript +++ b/SConscript @@ -1,3 +1,5 @@ +import sys + Import('env', 'envNanobind', 'nanobind_lib', 'common') @@ -26,9 +28,14 @@ visionipc = env.Library('visionipc', vipc_objects) vipc_libs = envNanobind["LIBS"] + [visionipc, msgq] + common -envNanobind.Program(f'{visionipc_dir.abspath}/visionipc_pyx.so', [f'{visionipc_dir.abspath}/visionipc_py.cc', nanobind_lib], +visionipc_python = envNanobind.Program(f'{visionipc_dir.abspath}/visionipc_pyx.so', [f'{visionipc_dir.abspath}/visionipc_py.cc', nanobind_lib], LIBS=vipc_libs) +# Generate typing information from the compiled modules for local checks and wheels. +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) diff --git a/lefthook.yml b/lefthook.yml index 708f9d42d..43d664a3b 100644 --- a/lefthook.yml +++ b/lefthook.yml @@ -17,8 +17,6 @@ test: run: ruff check . ty: run: ty check . - stubs: - run: python tools/generate_stubs.py --check codespell: 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 diff --git a/msgq/ipc_pyx.pyi b/msgq/ipc_pyx.pyi deleted file mode 100644 index 15e6cecb7..000000000 --- a/msgq/ipc_pyx.pyi +++ /dev/null @@ -1,77 +0,0 @@ -from collections.abc import Sequence - - -class IpcError(Exception): - pass - -class MultiplePublishersError(IpcError): - pass - -def toggle_fake_events(enabled: bool) -> None: ... - -def set_fake_prefix(prefix: str | bytes) -> None: ... - -def get_fake_prefix() -> bytes: ... - -def delete_fake_prefix() -> None: ... - -def wait_for_one_event(events: Sequence[Event], timeout: int = -1) -> int: ... - -class Event: - def __init__(self) -> None: ... - - def set(self) -> None: ... - - def clear(self) -> int: ... - - def wait(self, timeout: int = -1) -> None: ... - - def peek(self) -> bool: ... - - @property - def fd(self) -> int: ... - - @property - def ptr(self) -> int: ... - -class SocketEventHandle: - def __init__(self, endpoint: str | bytes, identifier: str | bytes, override: bool) -> None: ... - - @property - def enabled(self) -> bool: ... - - @enabled.setter - def enabled(self, arg: bool, /) -> None: ... - - @property - def recv_called_event(self) -> Event: ... - - @property - def recv_ready_event(self) -> Event: ... - -class SubSocket: - def __init__(self) -> None: ... - - def connect(self, endpoint: str | bytes, conflate: bool = False, segment_size: int = 0) -> None: ... - - def setTimeout(self, timeout: int) -> None: ... - - def receive(self, non_blocking: bool = False) -> bytes | None: ... - -class Poller: - def __init__(self) -> None: ... - - def registerSocket(self, socket: SubSocket) -> None: ... - - def poll(self, timeout: int) -> list[SubSocket]: ... - -class PubSocket: - def __init__(self) -> None: ... - - def connect(self, endpoint: str | bytes, segment_size: int = 0) -> None: ... - - def send(self, data: bytes) -> None: ... - - def all_readers_updated(self) -> bool: ... - - def wait_for_readers(self, timeout: float = 1.0, interval: float = 0.001) -> None: ... diff --git a/msgq/tests/test_bindings.py b/msgq/tests/test_bindings.py deleted file mode 100644 index c7f39c786..000000000 --- a/msgq/tests/test_bindings.py +++ /dev/null @@ -1,73 +0,0 @@ -import gc -import threading -import unittest - -import msgq -from msgq.ipc_pyx import Event -from msgq.tests.test_messaging import random_sock - - -class TestBindings(unittest.TestCase): - def test_bytes_endpoints_and_binary_messages(self): - endpoint = random_sock().encode() - publisher = msgq.PubSocket() - subscriber = msgq.SubSocket() - publisher.connect(endpoint) - subscriber.connect(endpoint) - publisher.send(b"\x00\xffpayload\x00") - self.assertEqual(subscriber.receive(), b"\x00\xffpayload\x00") - with self.assertRaises(TypeError): - getattr(publisher, "send")("text") - - def test_fake_prefix_bytes(self): - try: - for prefix in ("prefix", b"prefix"): - msgq.set_fake_prefix(prefix) - self.assertEqual(msgq.get_fake_prefix(), prefix.encode() if isinstance(prefix, str) else prefix) - finally: - msgq.delete_fake_prefix() - - def test_event_keeps_handle_alive(self): - event = msgq.fake_event_handle(random_sock()).recv_called_event - gc.collect() - event.set() - self.assertTrue(event.peek()) - event.clear() - self.assertFalse(event.peek()) - - def test_event_wait_releases_gil(self): - handle = msgq.fake_event_handle(random_sock()) - event = handle.recv_called_event - for wait in (lambda: event.wait(1), lambda: msgq.wait_for_one_event([event], 1)): - with self.subTest(wait=wait): - event.clear() - timer = threading.Timer(0.01, event.set) - timer.start() - try: - wait() - self.assertTrue(event.peek()) - finally: - timer.join() - - def test_poller_socket_lifetime(self): - endpoint = random_sock() - publisher = msgq.pub_sock(endpoint) - poller = msgq.Poller() - subscriber = msgq.sub_sock(endpoint, poller=poller) - del subscriber - gc.collect() - publisher.send(b"message") - sockets = poller.poll(1000) - self.assertEqual(len(sockets), 1) - del poller - gc.collect() - self.assertEqual(sockets[0].receive(), b"message") - - def test_reader_timeout(self): - publisher = msgq.pub_sock(random_sock()) - with self.assertRaisesRegex(TimeoutError, "subscriber did not connect"): - publisher.wait_for_readers(timeout=0) - - def test_invalid_event(self): - with self.assertRaisesRegex(RuntimeError, "valid file descriptor"): - Event().wait(0) diff --git a/msgq/tests/test_typing.py b/msgq/tests/test_typing.py deleted file mode 100644 index 4f48b011d..000000000 --- a/msgq/tests/test_typing.py +++ /dev/null @@ -1,43 +0,0 @@ -from typing import TYPE_CHECKING, assert_type - -if TYPE_CHECKING: - import msgq - from msgq.ipc_pyx import Event - from msgq.visionipc import VisionBuf, VisionIpcClient, VisionIpcServer, get_endpoint_name - - publisher = msgq.pub_sock("typing") - subscriber = msgq.sub_sock("typing") - poller = msgq.Poller() - poller.registerSocket(subscriber) - assert_type(publisher.send(b"data"), None) - assert_type(publisher.all_readers_updated(), bool) - assert_type(subscriber.receive(), bytes | None) - assert_type(poller.poll(0), list[msgq.SubSocket]) - assert_type(msgq.drain_sock_raw(subscriber), list[bytes]) - assert_type(msgq.get_fake_prefix(), bytes) - - handle = msgq.fake_event_handle("typing", identifier=b"prefix") - handle.enabled = True - assert_type(handle.enabled, bool) - assert_type(handle.recv_called_event, Event) - assert_type(handle.recv_ready_event.fd, int) - assert_type(msgq.wait_for_one_event([handle.recv_called_event], 0), int) - - server = VisionIpcServer("typing") - client = VisionIpcClient("typing", 0, False) - assert_type(get_endpoint_name("typing", 0), str) - assert_type(VisionIpcClient.available_streams("typing", False), set[int]) - assert_type(client.connect(False), bool) - assert_type(client.width, int | None) - assert_type(client.buffer_len, int | None) - assert_type(client.recv(), VisionBuf | None) - assert_type(client.frame_id, int) - assert_type(client.valid, bool) - frame = client.recv() - if frame is not None: - assert_type(frame.data, memoryview) - assert_type(frame.width, int) - assert_type(frame.frame_id, int) - server.send(0, frame.data) - server.send(0, b"frame") - server.send(0, bytearray(b"frame")) diff --git a/msgq/visionipc/visionipc_pyx.pyi b/msgq/visionipc/visionipc_pyx.pyi deleted file mode 100644 index 767383fe2..000000000 --- a/msgq/visionipc/visionipc_pyx.pyi +++ /dev/null @@ -1,79 +0,0 @@ -def get_endpoint_name(name: str | bytes, stream: int) -> str: ... - -class VisionBuf: - @property - def data(self) -> memoryview: ... - - @property - def width(self) -> int: ... - - @property - def height(self) -> int: ... - - @property - def stride(self) -> int: ... - - @property - def uv_offset(self) -> int: ... - - @property - def idx(self) -> int: ... - - @property - def fd(self) -> int: ... - - @property - def frame_id(self) -> int: ... - -class VisionIpcServer: - def __init__(self, name: str | bytes) -> None: ... - - def create_buffers(self, tp: int, num_buffers: int, width: int, height: int) -> None: ... - - def create_buffers_with_sizes(self, tp: int, num_buffers: int, width: int, height: int, size: int, stride: int, uv_offset: int) -> None: ... - - def send(self, tp: int, data: bytes | bytearray | memoryview, frame_id: int = 0, timestamp_sof: int = 0, timestamp_eof: int = 0) -> None: ... - - def start_listener(self) -> None: ... - -class VisionIpcClient: - def __init__(self, name: str | bytes, stream: int, conflate: bool) -> None: ... - - @property - def width(self) -> int | None: ... - - @property - def height(self) -> int | None: ... - - @property - def stride(self) -> int | None: ... - - @property - def uv_offset(self) -> int | None: ... - - @property - def buffer_len(self) -> int | None: ... - - @property - def num_buffers(self) -> int: ... - - @property - def frame_id(self) -> int: ... - - @property - def timestamp_sof(self) -> int: ... - - @property - def timestamp_eof(self) -> int: ... - - @property - def valid(self) -> bool: ... - - def recv(self, timeout_ms: int = 100) -> VisionBuf | None: ... - - def connect(self, blocking: bool) -> bool: ... - - def is_connected(self) -> bool: ... - - @staticmethod - def available_streams(name: str | bytes, block: bool) -> set[int]: ... diff --git a/setup.py b/setup.py index 45e2c172d..b7aa3986c 100644 --- a/setup.py +++ b/setup.py @@ -15,6 +15,8 @@ def build_extensions(self): target = Path(self.get_ext_fullpath(ext.name)) target.parent.mkdir(parents=True, exist_ok=True) self.copy_file(str(source), str(target)) + stub = source.with_suffix(".pyi") + self.copy_file(str(stub), str(target.parent / stub.name)) setup( diff --git a/test.sh b/test.sh index 2feab797a..dfe98c1ff 100755 --- a/test.sh +++ b/test.sh @@ -19,9 +19,7 @@ lefthook run test uv venv --python "$DIR/.venv/bin/python" "$TEST_DIR/.venv" uv pip install --python "$TEST_DIR/.venv/bin/python" "$DIR" cd "$TEST_DIR" - "$TEST_DIR/.venv/bin/python" -m unittest msgq.tests.test_messaging msgq.tests.test_bindings msgq.visionipc.tests.test_visionipc - cp "$DIR/msgq/tests/test_typing.py" "$TEST_DIR/check_types.py" - "$DIR/.venv/bin/ty" check --python "$TEST_DIR/.venv/bin/python" "$TEST_DIR/check_types.py" + "$TEST_DIR/.venv/bin/python" -m unittest msgq.tests.test_messaging msgq.visionipc.tests.test_visionipc ) # *** all done *** diff --git a/tools/generate_stubs.py b/tools/generate_stubs.py deleted file mode 100644 index 5d0f75867..000000000 --- a/tools/generate_stubs.py +++ /dev/null @@ -1,31 +0,0 @@ -import argparse -from pathlib import Path -import subprocess -import sys -import tempfile - -ROOT = Path(__file__).resolve().parents[1] -MODULES = ("msgq.ipc_pyx", "msgq.visionipc.visionipc_pyx") - - -def main() -> None: - parser = argparse.ArgumentParser(description="Generate nanobind stubs, or check that committed stubs are current.") - parser.add_argument("--check", action="store_true") - args = parser.parse_args() - - with tempfile.TemporaryDirectory() as directory: - for module in MODULES: - target = ROOT.joinpath(*module.split(".")).with_suffix(".pyi") - output = Path(directory) / target.name if args.check else target - subprocess.check_call([sys.executable, "-m", "nanobind.stubgen", "-i", str(ROOT), "-m", module, "-o", str(output), "-q"]) - # Older nanobind versions leave blank lines when a stub has no imports. - content = output.read_text().lstrip() - if args.check: - if content != target.read_text(): - sys.exit(f"{target.relative_to(ROOT)} is out of date; run python tools/generate_stubs.py") - else: - output.write_text(content) - - -if __name__ == "__main__": - main() From 98eb355de3bfb9f6fb32bbe5589b8622f193d0eb Mon Sep 17 00:00:00 2001 From: Adeeb Shihadeh Date: Sun, 25 Oct 2026 12:24:20 -0700 Subject: [PATCH 3/4] Fix macOS wheel deployment target --- pyproject.toml | 3 +++ 1 file changed, 3 insertions(+) diff --git a/pyproject.toml b/pyproject.toml index 2735b6fcb..115237922 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -48,6 +48,9 @@ namespaces = false build = "cp311-* cp312-* cp313-* cp314-*" test-command = "python -m unittest discover -s msgq" +[tool.cibuildwheel.macos.environment] +MACOSX_DEPLOYMENT_TARGET = "10.13" + # https://beta.ruff.rs/docs/configuration/#using-pyprojecttoml [tool.ruff] lint.select = ["E", "F", "W", "PIE", "C4", "ISC", "RUF100", "A"] From a3b98f660ff83b7f92cf04dac58a62dc4c1c96c7 Mon Sep 17 00:00:00 2001 From: Adeeb Shihadeh Date: Sun, 25 Oct 2026 12:57:54 -0700 Subject: [PATCH 4/4] Fix nanobind review findings --- SConscript | 15 +++++++++------ msgq/ipc_py.cc | 14 ++++++-------- setup.py | 6 ++++-- 3 files changed, 19 insertions(+), 16 deletions(-) diff --git a/SConscript b/SConscript index 141d2c972..a9e1af170 100644 --- a/SConscript +++ b/SConscript @@ -1,9 +1,11 @@ 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([ @@ -13,7 +15,7 @@ msgq_objects = env.SharedObject([ ]) msgq = env.Library('msgq', msgq_objects) # Keep the extension module names compatible with existing Python imports. -msgq_python = envNanobind.Program('msgq/ipc_pyx.so', ['msgq/ipc_py.cc', nanobind_lib], LIBS=envNanobind["LIBS"]+[msgq]+common) +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'] @@ -28,13 +30,14 @@ visionipc = env.Library('visionipc', vipc_objects) vipc_libs = envNanobind["LIBS"] + [visionipc, msgq] + common -visionipc_python = envNanobind.Program(f'{visionipc_dir.abspath}/visionipc_pyx.so', [f'{visionipc_dir.abspath}/visionipc_py.cc', nanobind_lib], +visionipc_python = envNanobind.Program(f'{visionipc_dir.abspath}/visionipc_pyx{python_suffix}', [f'{visionipc_dir.abspath}/visionipc_py.cc', nanobind_lib], LIBS=vipc_libs) -# Generate typing information from the compiled modules for local checks and wheels. -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"') +# 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) diff --git a/msgq/ipc_py.cc b/msgq/ipc_py.cc index 1c6e81a35..f0b785891 100644 --- a/msgq/ipc_py.cc +++ b/msgq/ipc_py.cc @@ -6,7 +6,6 @@ #include #include #include -#include #include #include @@ -105,13 +104,12 @@ NB_MODULE(ipc_pyx, m) { }, nb::arg("data")) .def("all_readers_updated", &PubSocket::all_readers_updated) .def("wait_for_readers", [](PubSocket &socket, double timeout, double interval) { - { - nb::gil_scoped_release release; - auto deadline = std::chrono::steady_clock::now() + std::chrono::duration(timeout); - while (std::chrono::steady_clock::now() < deadline) { - if (socket.all_readers_updated()) return; - std::this_thread::sleep_for(std::chrono::duration(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(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(); diff --git a/setup.py b/setup.py index b7aa3986c..4fa17fc0e 100644 --- a/setup.py +++ b/setup.py @@ -2,6 +2,7 @@ from pathlib import Path import subprocess import sys +import sysconfig from setuptools import Extension, setup from setuptools.command.build_ext import build_ext @@ -11,11 +12,12 @@ class SConsBuildExt(build_ext): def build_extensions(self): subprocess.check_call([sys.executable, "-m", "SCons", "--minimal", "-j", str(self.parallel or os.cpu_count() or 1)]) for ext in self.extensions: - source = Path(*ext.name.split(".")).with_suffix(".so") + module = Path(*ext.name.split(".")) + source = module.with_suffix(sysconfig.get_config_var("EXT_SUFFIX")) target = Path(self.get_ext_fullpath(ext.name)) target.parent.mkdir(parents=True, exist_ok=True) self.copy_file(str(source), str(target)) - stub = source.with_suffix(".pyi") + stub = module.with_suffix(".pyi") self.copy_file(str(stub), str(target.parent / stub.name))