Add not_foreign cast flag and nb_type_put() allow_foreign flag so we can prevent mutual recursion between two frameworks failing to perform a cast. Simplify enum destruction. Clean up some things I noticed while updating the pybind11 PR.

Author: oremanj

include/nanobind/nb_attr.h

@@ -207,7 +207,10 @@ enum cast_flags : uint8_t {
     // This implies that objects added to the cleanup list may be
     // released immediately after the caster's final output value is
     // obtained, i.e., before it is used.
-    manual = (1 << 3)
+    manual = (1 << 3),
+
+    // Disallow satisfying this cast with a foreign framework's binding
+    not_foreign = (1 << 4),
 };
 
 

include/nanobind/nb_class.h

@@ -205,9 +205,6 @@ enum class enum_flags : uint32_t {
 
     /// Is the underlying enumeration type Flag?
     is_flag                  = (1 << 3),
-
-    /// Was the enum successfully registered with nanobind?
-    is_registered            = (1 << 4),
 };
 
 struct enum_init_data {

include/nanobind/nb_lib.h

@@ -289,7 +289,8 @@ NB_CORE bool nb_type_get(const std::type_info *t, PyObject *o, uint8_t flags,
 /// Cast a C++ type instance into a Python object
 NB_CORE PyObject *nb_type_put(const std::type_info *cpp_type, void *value,
                               rv_policy rvp, cleanup_list *cleanup,
-                              bool *is_new = nullptr) noexcept;
+                              bool *is_new = nullptr,
+                              bool allow_foreign = true) noexcept;
 
 // Special version of nb_type_put for polymorphic classes
 NB_CORE PyObject *nb_type_put_p(const std::type_info *cpp_type,

include/nanobind/stl/unique_ptr.h

@@ -94,14 +94,12 @@ struct type_caster<std::unique_ptr<T, Deleter>> {
         // Stash source python object
         src = src_;
 
-        // Don't accept foreign types; they can't relinquish ownership
-        if (!src.is_none() && !inst_check(src))
-            return false;
-
-        /* Try casting to a pointer of the underlying type. We pass flags=0 and
-           cleanup=nullptr to prevent implicit type conversions (they are
-           problematic since the instance then wouldn't be owned by 'src') */
-        return caster.from_python(src_, 0, nullptr);
+        /* Try casting to a pointer of the underlying type. We pass
+           cleanup=nullptr and !(flags & convert) to prevent implicit type
+           conversions, which are problematic since the instance then wouldn't
+           be owned by 'src'. Also disable casting from a foreign type since it
+           wouldn't be able to relinquish ownership. */
+        return caster.from_python(src_, (uint8_t) cast_flags::not_foreign, nullptr);
     }
 
     template <typename T2>

src/nb_enum.cpp

@@ -73,8 +73,6 @@ PyObject *enum_create(enum_init_data *ed) noexcept {
         type_init_data *t = (type_init_data *) p;
         delete (enum_map *) t->enum_tbl.fwd;
         delete (enum_map *) t->enum_tbl.rev;
-        if (t->flags & (uint32_t) enum_flags::is_registered)
-            nb_type_unregister(t);
         free((char*) t->name);
         delete t;
     });
@@ -88,7 +86,12 @@ PyObject *enum_create(enum_init_data *ed) noexcept {
         return tp;
     }
 
-    t->flags |= (uint32_t) enum_flags::is_registered;
+    // Unregister the enum type when it begins being finalized
+    keep_alive(result.ptr(), t, [](void *p) noexcept {
+        nb_type_unregister((type_init_data *) p);
+    });
+
+    // Delete typeinfo only when the type's dict is cleared
     result.attr("__nb_enum__") = tie_lifetimes;
 
     make_immortal(result.ptr());
@@ -180,7 +183,7 @@ bool enum_from_python(const std::type_info *tp,
 
 #if !defined(NB_DISABLE_INTEROP)
     auto try_foreign = [=, &has_foreign]() -> bool {
-        if (has_foreign) {
+        if (has_foreign && !(flags & (uint8_t) cast_flags::not_foreign)) {
             void *ptr = nb_type_get_foreign(internals, tp, o, flags, cleanup);
             if (ptr) {
                 // Copy from the C++ enum object to our output integer.

src/nb_foreign.cpp

@@ -47,13 +47,14 @@ static void *nb_foreign_from_python(pymb_binding *binding,
                                                             PyObject *obj),
                                     void *keep_referenced_ctx) noexcept {
     cleanup_list cleanup{nullptr};
+    uint8_t flags = (uint8_t) cast_flags::not_foreign;
+    if (convert)
+        flags |= (uint8_t) cast_flags::convert;
     auto *td = (type_data *) binding->context;
     if (td->align == 0) { // enum
         int64_t value;
         if (keep_referenced &&
-            enum_from_python(td->type, pyobj, &value, td->size,
-                             convert ? uint8_t(cast_flags::convert) : 0,
-                             nullptr)) {
+            enum_from_python(td->type, pyobj, &value, td->size, flags, nullptr)) {
             bytes holder{(uint8_t *) &value + NB_BIG_ENDIAN * (8 - td->size),
                          td->size};
             keep_referenced(keep_referenced_ctx, holder.ptr());
@@ -63,8 +64,7 @@ static void *nb_foreign_from_python(pymb_binding *binding,
     }
 
     void *result = nullptr;
-    bool ok = nb_type_get(td->type, pyobj,
-                          convert ? uint8_t(cast_flags::convert) : 0,
+    bool ok = nb_type_get(td->type, pyobj, flags,
                           keep_referenced ? &cleanup : nullptr, &result);
     if (keep_referenced) {
         // Move temporary references from our `cleanup_list` to our caller's
@@ -109,7 +109,8 @@ static PyObject *nb_foreign_to_python(pymb_binding *binding,
         // unless a pyobject wrapper already exists.
         rvp = rv_policy::none;
     }
-    return nb_type_put(td->type, cobj, rvp, &cleanup, nullptr);
+    return nb_type_put(td->type, cobj, rvp, &cleanup,
+                       /* is_new */ nullptr, /* allow_foreign */ false);
 }
 
 static int nb_foreign_keep_alive(PyObject *nurse,
@@ -153,10 +154,12 @@ static int nb_foreign_translate_exception(void *eptr) noexcept {
         std::rethrow_exception(e);
     } catch (python_error &e) {
         e.restore();
+        return 1;
     } catch (builtin_exception &e) {
         if (!set_builtin_exception_status(e))
             PyErr_SetString(PyExc_SystemError, "foreign function threw "
                             "nanobind::next_overload()");
+        return 1;
     } catch (...) { e = std::current_exception(); }
     return 0;
 }
@@ -521,10 +524,10 @@ void *nb_type_try_foreign(nb_internals *internals_,
 #if defined(NB_FREE_THREADED)
     auto per_thread_guard = nb_type_lock_c2p_fast(internals_);
     nb_type_map_fast &type_c2p_fast = *per_thread_guard;
-    uint32_t updates_count = per_thread_guard.updates_count();
 #else
     nb_type_map_fast &type_c2p_fast = internals_->type_c2p_fast;
 #endif
+    uint32_t update_count = type_c2p_fast.update_count;
     do {
         // We assume nb_type_c2p already ran for this type, so that there's
         // no need to handle a cache miss here.
@@ -563,16 +566,16 @@ void *nb_type_try_foreign(nb_internals *internals_,
                     return result;
 
 #if defined(NB_FREE_THREADED)
-                // Re-acquire lock to continue iteration. If we missed an
-                // update while the lock was released, start our lookup over
-                // in case the update removed the node we're on.
+                // Re-acquire lock to continue iteration
                 per_thread_guard = nb_type_lock_c2p_fast(internals_);
-                if (per_thread_guard.updates_count() != updates_count) {
+#endif
+                // If we missed an update during attempt(), start our lookup
+                // over in case the update removed the node we're on.
+                if (type_c2p_fast.update_count != update_count) {
                     // Concurrent update occurred; retry
-                    updates_count = per_thread_guard.updates_count();
+                    update_count = type_c2p_fast.update_count;
                     break;
                 }
-#endif
             }
             current = current->next;
         }

src/nb_internals.h

@@ -331,6 +331,7 @@ struct nb_type_map_fast {
     /// then return a reference to the stored value, which the caller may
     /// modify.
     void*& lookup_or_set(const std::type_info *ti, void *dflt) {
+        ++update_count;
         return data.try_emplace((void *) ti, dflt).first.value();
     }
 
@@ -348,11 +349,23 @@ struct nb_type_map_fast {
         auto it = data.find((void *) ti);
         if (it != data.end()) {
             it.value() = value;
+            ++update_count;
             return true;
         }
         return false;
     }
 
+    /// Number of times the map has been modified. Used in nb_type_try_foreign()
+    /// to detect cases where attempting to use one foreign binding for a type
+    /// may have invalidated the iterator needed to advance to the next one.
+    uint32_t update_count = 0;
+
+#if defined(NB_FREE_THREADED)
+    /// Mutex used by `nb_type_map_per_thread`, stored here because it fits
+    /// in padding this way.
+    Py_Mutex mutex{};
+#endif
+
   private:
     // Use a generic ptr->ptr map to avoid needing another instantiation of
     // robin_map. Keys are const std::type_info*. See TYPE MAPPING above for
@@ -390,35 +403,30 @@ struct nb_type_map_per_thread {
         }
         ~guard() {
             if (parent)
-                PyMutex_Unlock(&parent->mutex);
+                PyMutex_Unlock(&parent->map.mutex);
         }
 
         nb_type_map_fast& operator*() const { return parent->map; }
         nb_type_map_fast* operator->() const { return &parent->map; }
 
-        uint32_t updates_count() const { return parent->updates; }
-        void note_updated() { ++parent->updates; }
-
       private:
         friend nb_type_map_per_thread;
         explicit guard(nb_type_map_per_thread &parent_) : parent(&parent_) {
-            PyMutex_Lock(&parent->mutex);
+            PyMutex_Lock(&parent->map.mutex);
         }
         nb_type_map_per_thread *parent = nullptr;
     };
     guard lock() { return guard{*this}; }
 
-    // Mutex protecting accesses to `updates` and `map`
-    PyMutex mutex{};
-
-    // The number of times `map` has been modified
-    uint32_t updates = 0;
-
     nb_internals &internals;
+
+  private:
+    // Access to the map is only possible via `guard`, which holds a lock
     nb_type_map_fast map;
 
+  public:
     // In order to access or modify `next`, you must hold the nb_internals mutex
-    // (this->mutex is not needed for iteration)
+    // (this->map.mutex is not needed for iteration)
     nb_type_map_per_thread *next = nullptr;
 };
 #endif

src/nb_type.cpp

@@ -472,12 +472,9 @@ void nb_type_update_c2p_fast(const std::type_info *type, void *value) noexcept {
         for (nb_type_map_per_thread *cache =
                  internals_->type_c2p_per_thread_head;
              cache; cache = cache->next) {
-            auto guard = cache->lock();
-            bool found = nb_type_update_cache(*guard, it_alias->first,
-                                              (nb_alias_seq *) it_alias->second,
-                                              value);
-            if (found)
-                guard.note_updated();
+            nb_type_update_cache(*cache->lock(), it_alias->first,
+                                 (nb_alias_seq *) it_alias->second,
+                                 value);
         }
         // We can't require that we found a match, because the type might
         // have been cached only by a thread that has since exited.
@@ -1621,6 +1618,9 @@ void *nb_type_get_foreign(nb_internals *internals_,
                           PyObject *src,
                           uint8_t flags,
                           cleanup_list *cleanup) noexcept {
+    if (flags & (uint8_t) cast_flags::not_foreign)
+        return nullptr;
+
     struct capture {
         PyObject *src;
         uint8_t flags;
@@ -1993,7 +1993,8 @@ PyObject *nb_type_put_foreign(nb_internals *internals_,
 PyObject *nb_type_put(const std::type_info *cpp_type,
                       void *value, rv_policy rvp,
                       cleanup_list *cleanup,
-                      bool *is_new) noexcept {
+                      bool *is_new,
+                      bool allow_foreign) noexcept {
     // Convert nullptr -> None
     if (!value) {
         Py_INCREF(Py_None);
@@ -2016,7 +2017,7 @@ PyObject *nb_type_put(const std::type_info *cpp_type,
 
 #if !defined(NB_DISABLE_INTEROP)
     auto try_foreign = [=, &has_foreign]() -> PyObject* {
-        if (has_foreign)
+        if (has_foreign && allow_foreign)
             return nb_type_put_foreign(internals_, cpp_type, nullptr, value,
                                        rvp, cleanup, is_new);
         return nullptr;

tests/test_inter_module.py

@@ -347,11 +347,12 @@ def repeatedly_attempt_conversions():
         transitions = limit
         thread.join()
 
-    # typical numbers from my machine: with limit=100, the test takes 6sec,
-    # and num_failed and num_successful are each several 10k's
+    # typical numbers from my machine: with limit=5000, the test takes a
+    # decent fraction of a second, and num_failed and num_successful are each
+    # several 10k's
     print(num_failed, num_successful)
     assert num_successful > 0
-    assert num_failed > 0 or not free_threaded
+    assert num_failed > 0
 
 
 def test12_multi_and_implicit(clean):
@@ -386,8 +387,10 @@ def test12_multi_and_implicit(clean):
     # Now add the other direction
     t1.import_all()
     tf.export_all()
-    assert t2.check_shared(sf_raw)
     assert t2.check_shared(sf_nb)
+    # Still need an explicit import for non-C++ type
+    t2.import_for_interop_explicit(tf.RawShared)
+    assert t2.check_shared(sf_raw)
 
     # Test normally passing these various objects
     for mod in (t2, tf):