Create a metaclass for each supplemental data type

Author: oremanj

docs/api_core.rst

@@ -2059,9 +2059,20 @@ declarations in generated :ref:`stubs <stubs>`,
 
 .. cpp:struct:: template <typename T> supplement
 
-   Indicate that ``sizeof(T)`` bytes of memory should be set aside to
-   store supplemental data in the type object. See :ref:`Supplemental
-   type data <supplement>` for more information.
+   Indicate that the type being created (not its instances) should contain
+   a subobject of type ``T``, which can later be accessed using
+   :cpp:func:`nb::type_supplement\<T\>() <type_supplement>`. nanobind will
+   create a separate metaclass per type ``T``. See
+   :ref:`Supplemental type data <supplement>` for more information.
+
+   .. cpp::function:: supplement&& inheritable() &&
+
+      By default, specifying a supplement for a type prevents that type from
+      being inherited, since there would be no general way to initialize the
+      supplemental data for the derived type. If you instead pass the binding
+      annotation ``nb::supplement<T>().inheritable()``, inheritance will be
+      allowed. Any derived type gets a new default-constructed copy of the
+      supplemental data, not a copy of the data stored in the base type.
 
 .. cpp:struct:: type_slots
 
@@ -2866,6 +2877,13 @@ Type objects
    accesses may crash the interpreter. Also refer to
    :cpp:class:`nb::supplement\<T\> <supplement>`.
 
+.. cpp:function:: template <typename T> bool type_has_supplement(handle h)
+
+   Assuming that `h` represents a bound type (see :cpp:func:`type_check`),
+   check whether it stores supplemental data of type ``T``. If this function
+   returns true, then it is valid to call :cpp:func:`nb::type_supplement\<T\>(h)
+   <type_supplement>`.
+
 .. cpp:function:: str type_name(handle h)
 
    Return the full (module-qualified) name of a type object as a Python string.

docs/changelog.rst

@@ -25,7 +25,20 @@ Upcoming version (TBA)
   long-standing inconvenience. (PR `#778
   <https://github.com/wjakob/nanobind/pull/778>`__).
 
-* ABI version 16.
+- nanobind now creates a separate metaclass for each supplemental data type
+  named in a :cpp:class:`nb::supplement\<T\>() <supplement>` class
+  binding annotation. (Previously, it created a separate metaclass for each
+  **size** of supplemental data.) This enables the use of supplemental data
+  types with non-trivial construction and destruction. You can also now query
+  whether a nanobind type uses a particular supplement type via the new
+  function :cpp:func:`nb::type_has_supplement\<T\>() <type_has_supplement>`,
+  and can implement certain metaclass customizations by defining a method
+  ``static void T::init_metaclass(PyTypeObject *metatype)``. Review the
+  documentation of :ref:`supplemental type data <supplement>` for more
+  information about the new capabilities. (PR `#972
+  <https://github.com/wjakob/nanobind/pull/972>`__).
+
+- ABI version 16.
 
 
 Version 2.5.0 (Feb 2, 2025)

docs/lowlevel.rst

@@ -238,11 +238,10 @@ Here is what this might look like in an implementation:
 
   struct MyTensorMetadata {
       bool stored_on_gpu;
-      // ..
-      // should be a POD (plain old data) type
+      // ...
   };
 
-  // Register a new type MyTensor, and reserve space for sizeof(MyTensorMedadata)
+  // Register a new type MyTensor, and reserve space for MyTensorMedadata
   nb::class_<MyTensor> cls(m, "MyTensor", nb::supplement<MyTensorMedadata>())
 
   /// Mutable reference to 'MyTensorMedadata' portion in Python type object
@@ -251,13 +250,17 @@ Here is what this might look like in an implementation:
 
 The :cpp:class:`nb::supplement\<T\>() <supplement>` annotation implicitly also
 passes :cpp:class:`nb::is_final() <is_final>` to ensure that type objects with
-supplemental data cannot be subclassed in Python.
-
-nanobind requires that the specified type ``T`` be trivially default
-constructible. It zero-initializes the supplement when the type is first
-created but does not perform any further custom initialization or destruction.
-You can fill the supplement with different contents following the type
-creation, e.g., using the placement new operator.
+supplemental data cannot be subclassed in Python. If you do wish to allow
+subclassing, write ``nb::supplement<T>().inheritable()`` instead. Any subclasses
+will get a new default-constructed copy of the supplemental ``T`` data, **not**
+a copy of their base class's data.
+
+nanobind requires that the specified type ``T`` be either constructible from
+``PyTypeObject*`` (the type that contains this supplement instance) or default
+constructible. If ``T`` has a trivial default constructor, it will be
+zero-initialized when the type is first created; otherwise, the appropriate
+constructor will be run. You can fill the supplement with different contents
+following the type creation.
 
 The contents of the supplemental data are not directly visible to Python's
 cyclic garbage collector, which creates challenges if you want to reference
@@ -268,6 +271,71 @@ name that begins with the symbol ``@``, then nanobind will prevent Python
 code from rebinding or deleting the attribute after it has been set, making
 the borrowed reference reasonably safe.
 
+Metaclass customizations
+^^^^^^^^^^^^^^^^^^^^^^^^
+
+nanobind internally creates a separate metaclass for each distinct
+supplement type ``T`` used in your program. It is possible to customize this
+metaclass in some limited ways, by providing a static member function of the
+supplement type: ``static void T::init_metaclass(PyTypeObject *mcls)``.
+If provided, this function will be invoked once after the metaclass has been
+created but before any types that use it have been created. It can customize
+the metaclass's behavior by assigning attributes, including ``__dunder__``
+methods.
+
+It is not possible to provide custom :ref:`type slots <typeslots>` for the
+metaclass, nor can the metaclass itself contain user-provided data (try
+defining properties backed by global variables instead). It is not possible
+to customize the metaclass's base or its type, because nanobind requires all
+nanobind instances to have the same meta-metaclass for quick identification.
+
+The following example shows how you could customize
+``__instancecheck__`` to obtain results similar to :py:class:`abc.ABC`\'s
+support for virtual base classes.
+
+.. code-block:: cpp
+
+   struct HasVirtualSubclasses {
+       static void init_metaclass(PyTypeObject *mcls) {
+          nb::cpp_function_def(
+                   [](nb::handle cls, nb::handle instance) {
+                       if (!nb::type_check(cls) ||
+                           !nb::type_has_supplement<HasVirtualSubclasses>(cls))
+                           return false;
+
+                       auto& supp = nb::type_supplement<HasVirtualSubclasses>(cls);
+                       return PyType_IsSubtype((PyTypeObject *) instance.type(),
+                                               (PyTypeObject *) cls.ptr()) ||
+                              nb::borrow<nb::set>(supp.subclasses).contains(
+                                      instance.type());
+                   },
+                   nb::is_method(), nb::scope(metatype),
+                   nb::name("__instancecheck__"));
+       }
+
+       explicit HasVirtualSubclasses(PyTypeObject *tp) {
+           nb::set subclasses_set;
+           nb::handle(tp).attr("@subclasses") = subclasses_set;
+           subclasses = subclasses_set;
+       }
+
+       void register(nb::handle tp) {
+           nb::borrow<nb::set>(subclasses).add(tp);
+       }
+
+       nb::handle subclasses;
+   };
+
+   struct Collection {};
+   struct Set {};
+
+   auto coll = nb::class_<Collection>(m, "Collection",
+                                      nb::supplement<HasVirtualSubclasses>());
+   auto set = nb::class_<Set>(m, "Set").def(nb::init<>());
+   nb::type_supplement<HasVirtualSubclasses>(coll).register(set);
+
+   // assert isinstance(m.Set(), m.Collection)
+
 .. _typeslots:
 
 Customizing type creation

docs/porting.rst

@@ -338,7 +338,12 @@ Removed features include:
   annotation was removed. However, the same behavior can be achieved by
   creating unnamed arguments; see the discussion in the section on
   :ref:`keyword-only arguments <kw_only>`.
-- ○ **Metaclasses**: creating types with custom metaclasses is unsupported.
+- ○ **Metaclasses**: Creating types with arbitrary custom metaclasses is
+  unsupported. However, many of the things you might want a metaclass for can
+  be accomplished using :ref:`supplemental type data <supplement>`. Each
+  supplement type is associated with its own metaclass, so you can add some
+  customizations after the metaclass has been created, but nanobind does not
+  allow customizing the metaclass's type, bases, or slots.
 - ○ **Module-local bindings**: support was removed (both for types and exceptions).
 - ○ **Custom allocation**: C++ classes with an overloaded or deleted ``operator
   new`` / ``operator delete`` are not supported.

include/nanobind/nb_attr.h

@@ -125,7 +125,14 @@ struct kw_only {};
 struct lock_self {};
 
 template <size_t /* Nurse */, size_t /* Patient */> struct keep_alive {};
-template <typename T> struct supplement {};
+template <typename T> struct supplement {
+    bool is_inheritable = false;
+
+    supplement&& inheritable() && {
+        is_inheritable = true;
+        return std::move(*this);
+    }
+};
 template <typename T> struct intrusive_ptr {
     intrusive_ptr(void (*set_self_py)(T *, PyObject *) noexcept)
         : set_self_py(set_self_py) { }

include/nanobind/nb_class.h

@@ -131,14 +131,65 @@ struct type_data {
 #endif
 };
 
+/// Information about a type T used as nb::supplement<T>()
+struct supplement_data {
+    // These aliases are required to work around a MSVC bug
+    using constructor_t = void (*)(void *supp, PyTypeObject *tp);
+    using destructor_t = void (*)(void *supp);
+    using init_metaclass_t = void (*)(PyTypeObject *metatype);
+
+    const std::type_info *type;
+    size_t size;
+    constructor_t construct; // nullptr if trivial
+    destructor_t destruct; // nullptr if trivial
+    init_metaclass_t init_metaclass; // nullptr if not defined
+
+    // Cache the metaclass object that nanobind has created to serve as
+    // the type of all nanobind types that use this supplement type.
+    // Due to this cache, the cost of looking up a given supplement type is
+    // paid once per module rather than once per nanobind type that uses it.
+    mutable PyTypeObject *metaclass_cache = nullptr;
+
+    // Link together all the supplement_data structures for the same type
+    // in different extension modules, so that all of their metaclass_cache
+    // members can be cleared when the metaclass is destroyed.
+    mutable const supplement_data *next_in_chain = nullptr;
+};
+
+template <typename T, typename = void>
+inline supplement_data::init_metaclass_t init_metaclass_for = nullptr;
+
+template <typename T>
+inline supplement_data::init_metaclass_t init_metaclass_for<
+        T, std::void_t<decltype(&T::init_metaclass)>> = &T::init_metaclass;
+
+template <typename T, typename = int>
+inline supplement_data::constructor_t construct_supplement_for =
+        std::is_trivially_default_constructible_v<T> ? nullptr :
+                +[](void *p, PyTypeObject *) { new (p) T{}; };
+
+template <typename T>
+inline supplement_data::constructor_t construct_supplement_for<
+        T, enable_if_t<std::is_constructible_v<T, PyTypeObject *>>> =
+                +[](void *p, PyTypeObject *tp) { new (p) T{tp}; };
+
+template <typename T>
+inline const supplement_data supplement_data_for = {
+        &typeid(T), sizeof(T),
+        construct_supplement_for<T>,
+        std::is_trivially_destructible_v<T> ? nullptr :
+                +[](void *p) { ((T *) p)->~T(); },
+        init_metaclass_for<T>,
+};
+
 /// Information about a type that is only relevant when it is being created
 struct type_init_data : type_data {
     PyObject *scope;
     const std::type_info *base;
     PyTypeObject *base_py;
     const char *doc;
     const PyType_Slot *type_slots;
-    size_t supplement;
+    const supplement_data *supplement;
 };
 
 NB_INLINE void type_extra_apply(type_init_data &t, const handle &h) {
@@ -184,13 +235,13 @@ NB_INLINE void type_extra_apply(type_init_data & t, const sig &s) {
 }
 
 template <typename T>
-NB_INLINE void type_extra_apply(type_init_data &t, supplement<T>) {
-    static_assert(std::is_trivially_default_constructible_v<T>,
-                  "The supplement must be a POD (plain old data) type");
+NB_INLINE void type_extra_apply(type_init_data &t, supplement<T> supp) {
     static_assert(alignof(T) <= alignof(void *),
                   "The alignment requirement of the supplement is too high.");
-    t.flags |= (uint32_t) type_init_flags::has_supplement | (uint32_t) type_flags::is_final;
-    t.supplement = sizeof(T);
+    t.flags |= (uint32_t) type_init_flags::has_supplement;
+    if (!supp.is_inheritable)
+        t.flags |= (uint32_t) type_flags::is_final;
+    t.supplement = &supplement_data_for<T>;
 }
 
 enum class enum_flags : uint32_t {
@@ -287,6 +338,8 @@ inline size_t type_align(handle h) { return detail::nb_type_align(h.ptr()); }
 inline const std::type_info& type_info(handle h) { return *detail::nb_type_info(h.ptr()); }
 template <typename T>
 inline T &type_supplement(handle h) { return *(T *) detail::nb_type_supplement(h.ptr()); }
+template <typename T>
+inline bool type_has_supplement(handle h) { return detail::nb_type_has_supplement(h.ptr(), &typeid(T)); }
 inline str type_name(handle h) { return steal<str>(detail::nb_type_name(h.ptr())); }
 
 // Low level access to nanobind instance objects

include/nanobind/nb_lib.h

@@ -316,6 +316,10 @@ NB_CORE void nb_type_restore_ownership(PyObject *o, bool cpp_delete) noexcept;
 /// Get a pointer to a user-defined 'extra' value associated with the nb_type t.
 NB_CORE void *nb_type_supplement(PyObject *t) noexcept;
 
+/// Determine whether the nb_type t has a supplement whose type matches supp_type.
+NB_CORE bool nb_type_has_supplement(PyObject *t,
+                                    const std::type_info *supp_type) noexcept;
+
 /// Check if the given python object represents a nanobind type
 NB_CORE bool nb_type_check(PyObject *t) noexcept;
 

src/nb_internals.cpp

@@ -31,6 +31,9 @@ extern int nb_bound_method_traverse(PyObject *, visitproc, void *);
 extern int nb_bound_method_clear(PyObject *);
 extern void nb_bound_method_dealloc(PyObject *);
 extern PyObject *nb_method_descr_get(PyObject *, PyObject *, PyObject *);
+extern PyTypeObject *nb_meta_new(nb_internals *internals,
+                                 const supplement_data *supplement) noexcept;
+extern void nb_meta_dealloc(PyObject *);
 
 #if PY_VERSION_HEX >= 0x03090000
 #  define NB_HAVE_VECTORCALL_PY39_OR_NEWER NB_HAVE_VECTORCALL
@@ -40,6 +43,7 @@ extern PyObject *nb_method_descr_get(PyObject *, PyObject *, PyObject *);
 
 static PyType_Slot nb_meta_slots[] = {
     { Py_tp_base, nullptr },
+    { Py_tp_dealloc, (void *) nb_meta_dealloc },
     { 0, nullptr }
 };
 
@@ -363,8 +367,12 @@ NB_NOINLINE void init(const char *name) {
     p->nb_module = PyModule_NewObject(nb_name.ptr());
 
     nb_meta_slots[0].pfunc = (PyObject *) &PyType_Type;
+#if PY_VERSION_HEX >= 0x030C0000
+    nb_meta_spec.basicsize = -(int) sizeof(const supplement_data*);
+#else
+    nb_meta_spec.basicsize = (int) (PyType_Type.tp_basicsize + sizeof(const supplement_data*));
+#endif
     nb_meta_cache = p->nb_meta = (PyTypeObject *) PyType_FromSpec(&nb_meta_spec);
-    p->nb_type_dict = PyDict_New();
     p->nb_func = (PyTypeObject *) PyType_FromSpec(&nb_func_spec);
     p->nb_method = (PyTypeObject *) PyType_FromSpec(&nb_method_spec);
     p->nb_bound_method = (PyTypeObject *) PyType_FromSpec(&nb_bound_method_spec);
@@ -379,7 +387,7 @@ NB_NOINLINE void init(const char *name) {
         p->shards[i].inst_c2p.min_load_factor(.1f);
     }
 
-    check(p->nb_module && p->nb_meta && p->nb_type_dict && p->nb_func &&
+    check(p->nb_module && p->nb_meta && p->nb_func &&
               p->nb_method && p->nb_bound_method,
           "nanobind::detail::init(): initialization failed!");
 
@@ -431,6 +439,8 @@ NB_NOINLINE void init(const char *name) {
     is_alive_ptr = &is_alive_value;
     p->is_alive_ptr = is_alive_ptr;
 
+    p->nb_type_0 = nb_meta_new(p, nullptr);
+
 #if PY_VERSION_HEX < 0x030C0000 && !defined(PYPY_VERSION)
     /* The implementation of typing.py on CPython <3.12 tends to introduce
        spurious reference leaks that upset nanobind's leak checker. The