Better handling for Python subclasses of types that use nb::new_() (#859)

* Better handling for Python subclasses of types that use nb::new_()

In Python, `__new__` takes a first argument indicating the type of object it should create, which allows one to derive from classes that use `__new__` and obtain reasonable semantics. nanobind's `nb::new_()` wrapper currently ignores this argument, with somewhat surprising results:
```
// C++
struct Base { ... };
nb::class_<Base>(m, "Base").def(nb::new_(...))...;

class Derived(mod.Base): pass

>>> type(Derived()) is mod.Base
True
```

This PR makes that case work more like it does in Python, so that `Derived()` produces an instance of `Derived`. This is possible safely because both `Base` and `Derived` believe they wrap the same C++ type. Since we can't easily intervene before the `Base` object is created, we use a call policy to intervene afterwards, and return an object of type `Derived` that wraps a `Base*` without ownership. The original `Base` object's lifetime is maintained via a keep-alive from the returned wrapper. There is not much actual code, but it's subtle so the comments are pretty long.

This feature requires allowing a call policy's postcall hook to modify the return value, which was not previously supported. I chose to implement this by allowing the postcall hook to take the return value via `PyObject*&`; then a passed `PyObject*` can initialize either that or to the documented `handle`. I didn't document this new capability, because it's somewhat obscure and easy to mess up the reference counting; I figure anyone that really needs it will be able to figure it out. An alternative if you don't want to add this ability to call policies in general would be to define a new function attribute just for the new-fixup case, but that felt more invasive to me.

Author: oremanj

docs/changelog.rst

@@ -56,6 +56,15 @@ Version TBD (not yet released)
   <https://github.com/wjakob/nanobind/pull/847>`__, commit `b95eb7
   <https://github.com/wjakob/nanobind/commit/b95eb755b5a651a40562002be9ca8a4c6bf0acb9>`__).
 
+- It is now possible to create Python subclasses of C++ classes that
+  define their constructor bindings using :cpp:struct:`nb::new_() <new_>`.
+  Previously, attempting to instantiate such a Python subclass would instead
+  produce an instance of the base C++ type. Note that it is still not possible
+  to override virtual methods in such a Python subclass, because the object
+  returned by the :cpp:struct:`new_() <new_>` constructor will generally
+  not be an instance of the alias/trampoline type.
+  (PR `#859 <https://github.com/wjakob/nanobind/pull/859>`__)
+
 Version 2.4.0 (Dec 6, 2024)
 ---------------------------
 

docs/classes.rst

@@ -1104,6 +1104,29 @@ type object that Python passes as the first argument of ``__new__``
 ``.def_static("__new__", ...)`` and matching ``.def("__init__", ...)``
 yourself.
 
+Two limitations of :cpp:struct:`nb::new_ <new_>` are worth noting:
+
+* The possibilities for Python-side inheritance from C++ classes that
+  are bound using :cpp:struct:`nb::new_ <new_>` constructors are substantially
+  reduced. Simple inheritance situations (``class PyPet(Pet): ...``) should
+  work OK, but you can't :ref:`override virtual functions <trampolines>`
+  in Python (because the C++ object returned by :cpp:struct:`new_ <new_>`
+  doesn't contain the Python trampoline glue), and if :cpp:struct:`new_ <new_>`
+  is used to implement a polymorphic factory (like if ``Pet::make()`` could
+  return an instance of ``Cat``) then Python-side inheritance won't work at all.
+
+* A given C++ class must expose all of its constructors via ``__new__`` or
+  all via ``__init__``, rather than a mixture of the two.
+  The only case where a class should bind both of these methods is
+  if the ``__init__`` methods are all stubs that do nothing.
+  This is because nanobind internally optimizes object instantiation by
+  caching the method that should be used for constructing instances of each
+  given type, and that optimization doesn't support trying both methods.
+  If you really need to combine nontrivial ``__new__`` and nontrivial
+  ``__init__`` in the same type, you can disable the optimization by
+  defining a :ref:`custom type slot <typeslots>` of ``Py_tp_new`` or
+  ``Py_tp_init``.
+
 .. note::
 
    Unpickling an object of type ``Foo`` normally requires that

docs/faq.rst

@@ -379,7 +379,7 @@ Stable ABI extensions are convenient because they can be reused across Python
 versions, but this unfortunately only works on Python 3.12 and newer. Nanobind
 crucially depends on several `features
 <https://docs.python.org/3/whatsnew/3.12.html#c-api-changes>`__ that were added
-in version 3.12 (specifically, `PyType_FromMetaclass()`` and limited API
+in version 3.12 (specifically, ``PyType_FromMetaclass()`` and limited API
 bindings of the vector call protocol).
 
 Policy on Clang-Tidy, ``-Wpedantic``, etc.

include/nanobind/nb_attr.h

@@ -459,11 +459,11 @@ process_postcall(PyObject **args, std::integral_constant<size_t, NArgs>,
 template <size_t NArgs, typename Policy>
 NB_INLINE void
 process_postcall(PyObject **args, std::integral_constant<size_t, NArgs> nargs,
-                 PyObject *result, call_policy<Policy> *) {
+                 PyObject *&result, call_policy<Policy> *) {
     // result_guard avoids leaking a reference to the return object
     // if postcall throws an exception
     object result_guard = steal(result);
-    Policy::postcall(args, nargs, handle(result));
+    Policy::postcall(args, nargs, result);
     result_guard.release();
 }
 

include/nanobind/nb_class.h

@@ -416,6 +416,39 @@ namespace detail {
             }
         }
     }
+
+    // Call policy that ensures __new__ returns an instance of the correct
+    // Python type, even when deriving from the C++ class in Python
+    struct new_returntype_fixup_policy {
+        static inline void precall(PyObject **, size_t,
+                                   detail::cleanup_list *) {}
+        NB_NOINLINE static inline void postcall(PyObject **args, size_t,
+                                                PyObject *&ret) {
+            handle type_requested = args[0];
+            if (ret == nullptr || !type_requested.is_type())
+                return; // somethign strange about this call; don't meddle
+            handle type_created = Py_TYPE(ret);
+            if (type_created.is(type_requested))
+                return; // already created the requested type so no fixup needed
+
+            if (type_check(type_created) &&
+                PyType_IsSubtype((PyTypeObject *) type_requested.ptr(),
+                                 (PyTypeObject *) type_created.ptr()) &&
+                type_info(type_created) == type_info(type_requested)) {
+                // The new_ constructor returned an instance of a bound type T.
+                // The user wanted an instance of some python subclass S of T.
+                // Since both wrap the same C++ type, we can satisfy the request
+                // by returning a pyobject of type S that wraps a C++ T*, and
+                // handling the lifetimes by having that pyobject keep the
+                // already-created T pyobject alive.
+                object wrapper = inst_reference(type_requested,
+                                                inst_ptr<void>(ret),
+                                                /* parent = */ ret);
+                handle(ret).dec_ref();
+                ret = wrapper.release().ptr();
+            }
+        }
+    };
 }
 
 template <typename Func, typename Sig = detail::function_signature_t<Func>>
@@ -446,13 +479,15 @@ struct new_<Func, Return(Args...)> {
             return func_((detail::forward_t<Args>) args...);
         };
 
+        auto policy = call_policy<detail::new_returntype_fixup_policy>();
         if constexpr ((std::is_base_of_v<arg, Extra> || ...)) {
             // If any argument annotations are specified, add another for the
             // extra class argument that we don't forward to Func, so visible
             // arg() annotations stay aligned with visible function arguments.
-            cl.def_static("__new__", std::move(wrapper), arg("cls"), extra...);
+            cl.def_static("__new__", std::move(wrapper), arg("cls"), extra...,
+                          policy);
         } else {
-            cl.def_static("__new__", std::move(wrapper), extra...);
+            cl.def_static("__new__", std::move(wrapper), extra..., policy);
         }
         cl.def("__init__", [](handle, Args...) {}, extra...);
     }

tests/test_classes.py

@@ -901,6 +901,27 @@ def test46_custom_new():
     assert t.NewStar("hi", "lo", value=10).value == 12
     assert t.NewStar(value=10, other="blah").value == 20
 
+    # Make sure a Python class that derives from a C++ class that uses
+    # nb::new_() can be instantiated producing the correct Python type
+    class FancyInt(t.UniqueInt):
+        @staticmethod
+        def the_answer():
+            return 42
+
+        @property
+        def value_as_string(self):
+            return str(self.value())
+
+    f1 = FancyInt(10)
+    f2 = FancyInt(20)
+    # The derived-type wrapping doesn't preserve Python identity...
+    assert f1 is not FancyInt(10)
+    # ... but does preserve C++ identity
+    assert f1.lookups() == u4.lookups() == 3  # u4, f1, and anonymous
+    assert f1.the_answer() == f2.the_answer() == 42
+    assert f1.value_as_string == "10"
+    assert f2.value_as_string == "20"
+
 def test47_inconstructible():
     with pytest.raises(TypeError, match="no constructor defined"):
         t.Foo()