From 34d0be1bbeb54b8265456fd3a4a50e98f93fe6d4 Mon Sep 17 00:00:00 2001 From: Wenzel Jakob Date: Fri, 21 Oct 2022 15:08:23 +0200 Subject: [PATCH] nb::make_iterator support --- README.md | 23 ++++-- include/nanobind/make_iterator.h | 133 +++++++++++++++++++++++++++++++ src/nb_internals.cpp | 3 +- src/nb_internals.h | 3 + src/nb_type.cpp | 25 +++--- tests/CMakeLists.txt | 2 + tests/test_make_iterator.cpp | 42 ++++++++++ tests/test_make_iterator.py | 35 ++++++++ 8 files changed, 250 insertions(+), 16 deletions(-) create mode 100644 include/nanobind/make_iterator.h create mode 100644 tests/test_make_iterator.cpp create mode 100644 tests/test_make_iterator.py diff --git a/README.md b/README.md index d14da074..6d12d205 100644 --- a/README.md +++ b/README.md @@ -393,15 +393,17 @@ changes are detailed below. keeps track of Python temporaries that are created by the conversion, and which need to be deallocated after a function call has taken place. `flags` and `cleanup` should be passed to any recursive usage of - `type_caster::from_python()`. + `type_caster::from_python()`. If casting fails due to a Python exception, + the function should clear it (`PyErr_Clear()`) and return false. If a + severe error condition arises that should be reported, use Python warning + API calls for this, e.g. `PyErr_WarnFormat()`. - `cast()` was renamed to `from_cpp()`. The function takes a return value - policy (as before) and a `cleanup_list *` pointer. + policy (as before) and a `cleanup_list *` pointer. If casting fails due to + a Python exception, the function should *leave the error set* (note the + asymmetry compared to ``from_python()``) and return ``nullptr``. - Both functions must be marked as `noexcept`. In contrast to _pybind11_, - errors during type casting are only propagated using status codes. If a - severe error condition arises that should be reported, use Python warning API - calls for this, e.g. `PyErr_WarnFormat()`. + Both functions must be marked as `noexcept`. Note that the cleanup list is only available when `from_python()` or `from_cpp()` are called as part of function dispatch, while usage by @@ -412,6 +414,15 @@ changes are detailed below. caster](https://github.com/wjakob/nanobind/blob/master/include/nanobind/stl/pair.h) may be useful as a reference for these changes. +- Use of the ``nb::make_iterator()``, ``nb::make_key_iterator()``, and + ``nb::make_value_iterator()`` requires including the additional header file + ``nanobind/make_iterator.h``. The interface of these functions has also + slightly changed: all take a Python scope and a name as first and second + arguments, which are used to permanently "install" the iterator type (which + is created on demand). See the The [test + suite](https://github.com/wjakob/nanobind/blob/master/tests/test_make_iterator.cpp) + for a worked out example. + - The following types and functions were renamed: | _pybind11_ | _nanobind_ | diff --git a/include/nanobind/make_iterator.h b/include/nanobind/make_iterator.h new file mode 100644 index 00000000..dd16881f --- /dev/null +++ b/include/nanobind/make_iterator.h @@ -0,0 +1,133 @@ +/* + nanobind/make_iterator.h: nb::make_[key,value_]iterator() + + This implementation is a port from pybind11 with minimal adjustments. + + All rights reserved. Use of this source code is governed by a + BSD-style license that can be found in the LICENSE file. +*/ + +#pragma once + +#include +#include + +NAMESPACE_BEGIN(NB_NAMESPACE) +NAMESPACE_BEGIN(detail) + +/* There are a large number of apparently unused template arguments because + each combination requires a separate nb::class_ registration. */ +template +struct iterator_state { + Iterator it; + Sentinel end; + bool first_or_done; +}; + +// Note: these helpers take the iterator by non-const reference because some +// iterators in the wild can't be dereferenced when const. +template struct iterator_access { + using result_type = decltype(*std::declval()); + result_type operator()(Iterator &it) const { return *it; } +}; + +template struct iterator_key_access { + using result_type = decltype((*std::declval()).first); + result_type operator()(Iterator &it) const { return (*it).first; } +}; + +template struct iterator_value_access { + using result_type = decltype((*std::declval()).second); + result_type operator()(Iterator &it) const { return (*it).second; } +}; + +template +iterator make_iterator_impl(handle scope, const char *name, + Iterator &&first, Sentinel &&last, + Extra &&...extra) { + using State = iterator_state; + + if (!type().is_valid()) { + class_(scope, name) + .def("__iter__", [](handle h) { return h; }) + .def("__next__", + [](State &s) -> ValueType { + if (!s.first_or_done) + ++s.it; + else + s.first_or_done = false; + + if (s.it == s.end) { + s.first_or_done = true; + throw stop_iteration(); + } + + return Access()(s.it); + }, + std::forward(extra)..., + Policy); + } + + return borrow(cast(State{ std::forward(first), + std::forward(last), true })); +} + +NAMESPACE_END(detail) + +/// Makes a python iterator from a first and past-the-end C++ InputIterator. +template ::result_type, + typename... Extra> +iterator make_iterator(handle scope, const char *name, Iterator &&first, Sentinel &&last, Extra &&...extra) { + return detail::make_iterator_impl, Policy, + Iterator, Sentinel, ValueType, Extra...>( + scope, name, std::forward(first), + std::forward(last), std::forward(extra)...); +} + +/// Makes an iterator over the keys (`.first`) of a iterator over pairs from a +/// first and past-the-end InputIterator. +template ::result_type, + typename... Extra> +iterator make_key_iterator(handle scope, const char *name, Iterator &&first, + Sentinel &&last, Extra &&...extra) { + return detail::make_iterator_impl, + Policy, Iterator, Sentinel, KeyType, + Extra...>( + scope, name, std::forward(first), + std::forward(last), std::forward(extra)...); +} + +/// Makes an iterator over the values (`.second`) of a iterator over pairs from a +/// first and past-the-end InputIterator. +template ::result_type, + typename... Extra> +iterator make_value_iterator(handle scope, const char *name, Iterator &&first, Sentinel &&last, Extra &&...extra) { + return detail::make_iterator_impl, + Policy, Iterator, Sentinel, ValueType, + Extra...>( + scope, name, std::forward(first), + std::forward(last), std::forward(extra)...); +} + +/// Makes an iterator over values of a container supporting `std::begin()`/`std::end()` +template +iterator make_iterator(handle scope, const char *name, Type &value, Extra &&...extra) { + return make_iterator(scope, name, std::begin(value), + std::end(value), + std::forward(extra)...); +} + +NAMESPACE_END(NB_NAMESPACE) diff --git a/src/nb_internals.cpp b/src/nb_internals.cpp index e1a32cf9..28d60ab8 100644 --- a/src/nb_internals.cpp +++ b/src/nb_internals.cpp @@ -17,7 +17,7 @@ /// Tracks the ABI of nanobind #ifndef NB_INTERNALS_VERSION -# define NB_INTERNALS_VERSION 3 +# define NB_INTERNALS_VERSION 4 #endif /// On MSVC, debug and release builds are not ABI-compatible! @@ -364,6 +364,7 @@ static void internals_make() { if (rv || !capsule || !nb_module) fail("nanobind::detail::internals_make(): allocation failed!"); Py_DECREF(capsule); + internals_p->nb_module = nb_module; internals_p->type_basicsize = cast(handle(&PyType_Type).attr("__basicsize__")); diff --git a/src/nb_internals.h b/src/nb_internals.h index 1fd5f9e5..76bf23d0 100644 --- a/src/nb_internals.h +++ b/src/nb_internals.h @@ -170,6 +170,9 @@ using keep_alive_set = py_set; struct nb_internals { + /// Internal nanobind module + PyObject *nb_module; + /// Registered metaclasses for nanobind classes and enumerations PyTypeObject *nb_type, *nb_enum; diff --git a/src/nb_type.cpp b/src/nb_type.cpp index e250d9da..ef2e8da2 100644 --- a/src/nb_type.cpp +++ b/src/nb_type.cpp @@ -316,16 +316,20 @@ PyObject *nb_type_new(const type_data *t) noexcept { PyObject *mod = nullptr; if (has_scope) { - if (PyModule_Check(t->scope)) { - mod = t->scope; - modname = getattr(t->scope, "__name__", handle()); - } else { - modname = getattr(t->scope, "__module__", handle()); + has_scope = t->scope != nullptr; + + if (has_scope) { + if (PyModule_Check(t->scope)) { + mod = t->scope; + modname = getattr(t->scope, "__name__", handle()); + } else { + modname = getattr(t->scope, "__module__", handle()); - object scope_qualname = getattr(t->scope, "__qualname__", handle()); - if (scope_qualname.is_valid()) - qualname = steal( - PyUnicode_FromFormat("%U.%U", scope_qualname.ptr(), name.ptr())); + object scope_qualname = getattr(t->scope, "__qualname__", handle()); + if (scope_qualname.is_valid()) + qualname = steal( + PyUnicode_FromFormat("%U.%U", scope_qualname.ptr(), name.ptr())); + } } } @@ -500,6 +504,9 @@ PyObject *nb_type_new(const type_data *t) noexcept { type_data *to = nb_type_data((PyTypeObject *) result); *to = *t; + if (!has_scope) + to->flags &= ~(uint32_t) type_flags::has_scope; + if (!intrusive_ptr && tb && (tb->flags & (uint32_t) type_flags::intrusive_ptr)) { to->flags |= (uint32_t) type_flags::intrusive_ptr; diff --git a/tests/CMakeLists.txt b/tests/CMakeLists.txt index 23cddcd3..5db50b9b 100644 --- a/tests/CMakeLists.txt +++ b/tests/CMakeLists.txt @@ -6,6 +6,7 @@ nanobind_add_module(test_enum_ext test_enum.cpp) nanobind_add_module(test_tensor_ext test_tensor.cpp) nanobind_add_module(test_intrusive_ext test_intrusive.cpp object.cpp object.h) nanobind_add_module(test_exception_ext test_exception.cpp object.cpp object.h) +nanobind_add_module(test_make_iterator_ext test_make_iterator.cpp) set(TEST_FILES test_functions.py @@ -16,6 +17,7 @@ set(TEST_FILES test_tensor.py test_intrusive.py test_exception.py + test_make_iterator.py ) if (NOT (CMAKE_CURRENT_SOURCE_DIR STREQUAL CMAKE_CURRENT_BINARY_DIR) OR MSVC) diff --git a/tests/test_make_iterator.cpp b/tests/test_make_iterator.cpp new file mode 100644 index 00000000..cd964ab2 --- /dev/null +++ b/tests/test_make_iterator.cpp @@ -0,0 +1,42 @@ +#include +#include +#include + +namespace nb = nanobind; + +NB_MODULE(test_make_iterator_ext, m) { + struct StringMap { + std::unordered_map map; + decltype(map.cbegin()) begin() const { return map.cbegin(); } + decltype(map.cend()) end() const { return map.cend(); } + }; + + nb::class_(m, "StringMap") + .def(nb::init<>()) + .def(nb::init>()) + .def("__iter__", + [](const StringMap &map) { + return nb::make_key_iterator(nb::type(), + "key_iterator", + map.begin(), + map.end()); + }, nb::keep_alive<0, 1>()) + .def("items", + [](const StringMap &map) { + return nb::make_iterator(nb::type(), + "item_iterator", + map.begin(), + map.end()); + }, nb::keep_alive<0, 1>()) + .def("values", [](const StringMap &map) { + return nb::make_value_iterator(nb::type(), + "value_iterator", + map.begin(), + map.end()); + }, nb::keep_alive<0, 1>()); + + nb::handle mod = m; + m.def("iterator_passthrough", [mod](nb::iterator s) -> nb::iterator { + return nb::make_iterator(mod, "pt_iterator", std::begin(s), std::end(s)); + }); +} diff --git a/tests/test_make_iterator.py b/tests/test_make_iterator.py new file mode 100644 index 00000000..5d609afa --- /dev/null +++ b/tests/test_make_iterator.py @@ -0,0 +1,35 @@ +import test_make_iterator_ext as t +import pytest + +data = [ + {}, + { 'a' : 'b' }, + { str(i) : chr(i) for i in range(1000) } +] + + +def test01_key_iterator(): + for d in data: + m = t.StringMap(d) + assert sorted(list(m)) == sorted(list(d)) + + +def test02_value_iterator(): + types = [] + for d in data: + m = t.StringMap(d) + types.append(type(m.values())) + assert sorted(list(m.values())) == sorted(list(d.values())) + assert types[0] is types[1] and types[1] is types[2] + + +def test03_items_iterator(): + for d in data: + m = t.StringMap(d) + assert sorted(list(m.items())) == sorted(list(d.items())) + + +def test04_passthrough_iterator(): + for d in data: + m = t.StringMap(d) + assert list(t.iterator_passthrough(m.values())) == list(m.values())