Skip to content

Commit

Permalink
Make nb::dict iterator non-copyable
Browse files Browse the repository at this point in the history
This commit changes the ``nb::dict`` iterator so that nanobind can
implement the recommendation from

https://docs.python.org/3.14/howto/free-threading-extensions.html#pydict-next

The primary goal of ``nb::internal::dict_iterator`` was to be able to write

```cpp
nb::dict my_dict = /* ... */;
for (auto [k, v] : my_dict) {
    // ....
}
```

This in fact the only associated feature that is explicitly mentioned in
the documentation, and this continues to work.

However, some undocumented features are lost:

- The dictionary iterator is no longer copyable. This is because it
  must acquire an exclusive lock to the underlying dictionary.

- The pre-increment operator ``++dict_it`` (which relied on copying) is
  gone. Post-increment continues to work, and that is enough for the
  loop structure mentioned above.
  • Loading branch information
wjakob committed Sep 20, 2024
1 parent 1de49f4 commit d9b2b21
Show file tree
Hide file tree
Showing 7 changed files with 32 additions and 18 deletions.
4 changes: 4 additions & 0 deletions docs/api_core.rst
Original file line number Diff line number Diff line change
Expand Up @@ -790,6 +790,10 @@ Wrapper classes
Return an item iterator that returns ``std::pair<handle, handle>``
key-value pairs analogous to ``iter(dict.items())`` in Python.

In free-threaded Python, the :cpp:class:``detail::dict_iterator`` class
acquires a lock to the underlying dictionary to enable the use of the
efficient but thread-unsafe ``PyDict_Next()`` Python C traversal routine.

.. cpp:function:: detail::dict_iterator end() const

Return a sentinel that ends the iteration.
Expand Down
4 changes: 2 additions & 2 deletions docs/free_threaded.rst
Original file line number Diff line number Diff line change
Expand Up @@ -235,8 +235,8 @@ necessary based on future experience and changes in Python itself.
Wrappers
________

:ref:`Wrapper types <wrappers>` like :cpp:class:`nb::list <list>` may be used in
multi-threaded code. Operations like :cpp:func:`nb::list::append()
:ref:`Wrapper types <wrappers>` like :cpp:class:`nb::list <list>` may be used
in multi-threaded code. Operations like :cpp:func:`nb::list::append()
<list::append>` internally acquire locks and behave just like their ordinary
Python counterparts. This means that race conditions can still occur without
larger-scale synchronization, but such races won't jeopardize the memory safety
Expand Down
2 changes: 1 addition & 1 deletion include/nanobind/nanobind.h
Original file line number Diff line number Diff line change
Expand Up @@ -48,10 +48,10 @@
#include "nb_error.h"
#include "nb_attr.h"
#include "nb_cast.h"
#include "nb_misc.h"
#include "nb_call.h"
#include "nb_func.h"
#include "nb_class.h"
#include "nb_misc.h"

#if defined(_MSC_VER)
# pragma warning(pop)
Expand Down
2 changes: 1 addition & 1 deletion include/nanobind/nb_call.h
Original file line number Diff line number Diff line change
Expand Up @@ -65,7 +65,7 @@ NB_INLINE void call_init(PyObject **args, PyObject *kwnames, size_t &nargs,
} else if constexpr (std::is_same_v<D, kwargs_proxy>) {
PyObject *key, *entry;
Py_ssize_t pos = 0;

ft_object_guard guard(value);
while (PyDict_Next(value.ptr(), &pos, &key, &entry)) {
Py_INCREF(key); Py_INCREF(entry);
args[kwargs_offset + nkwargs] = entry;
Expand Down
5 changes: 5 additions & 0 deletions include/nanobind/nb_defs.h
Original file line number Diff line number Diff line change
Expand Up @@ -173,6 +173,11 @@
# define NB_TYPE_GET_SLOT_IMPL 1
#endif

#define NB_NONCOPYABLE(X) \
X(const X &) = delete; \
X &operator=(const X &) = delete;


#define NB_MODULE_IMPL(name) \
extern "C" [[maybe_unused]] NB_EXPORT PyObject *PyInit_##name(); \
extern "C" NB_EXPORT PyObject *PyInit_##name()
Expand Down
4 changes: 0 additions & 4 deletions include/nanobind/nb_misc.h
Original file line number Diff line number Diff line change
Expand Up @@ -9,10 +9,6 @@

NAMESPACE_BEGIN(NB_NAMESPACE)

#define NB_NONCOPYABLE(X) \
X(const X &) = delete; \
X &operator=(const X &) = delete;

struct gil_scoped_acquire {
public:
NB_NONCOPYABLE(gil_scoped_acquire)
Expand Down
29 changes: 19 additions & 10 deletions include/nanobind/nb_types.h
Original file line number Diff line number Diff line change
Expand Up @@ -819,24 +819,29 @@ struct fast_iterator {

class dict_iterator {
public:
NB_NONCOPYABLE(dict_iterator)

using value_type = std::pair<handle, handle>;
using reference = const value_type;

dict_iterator() : h(), pos(-1) { }

dict_iterator() = default;
dict_iterator(handle h) : h(h), pos(0) {
#if defined(NB_FREE_THREADED)
PyCriticalSection_Begin(&cs, h.ptr());
#endif
increment();
}

dict_iterator& operator++() {
increment();
return *this;
#if defined(NB_FREE_THREADED)
~dict_iterator() {
if (h.ptr())
PyCriticalSection_End(&cs);
}
#endif

dict_iterator operator++(int) {
dict_iterator rv = *this;
dict_iterator& operator++() {
increment();
return rv;
return *this;
}

void increment() {
Expand All @@ -851,8 +856,12 @@ class dict_iterator {

private:
handle h;
Py_ssize_t pos;
PyObject *key = nullptr, *value = nullptr;
Py_ssize_t pos = -1;
PyObject *key = nullptr;
PyObject *value = nullptr;
#if defined(NB_FREE_THREADED)
PyCriticalSection cs { };
#endif
};

NB_IMPL_COMP(equal, Py_EQ)
Expand Down

0 comments on commit d9b2b21

Please sign in to comment.