Add custom_type_setup attribute

This allows for custom modifications to the PyHeapTypeObject prior to
calling `PyType_Ready`.  This may be used, for example, to define
`tp_traverse` and `tp_clear` functions.

Author: jbms

CMakeLists.txt

@@ -106,6 +106,7 @@ set(PYBIND11_HEADERS
     include/pybind11/detail/class.h
     include/pybind11/detail/common.h
     include/pybind11/detail/descr.h
+    include/pybind11/detail/function_view.h
     include/pybind11/detail/init.h
     include/pybind11/detail/internals.h
     include/pybind11/detail/type_caster_base.h

docs/advanced/classes.rst

@@ -1247,7 +1247,7 @@ Accessing the type object
 
 You can get the type object from a C++ class that has already been registered using:
 
-.. code-block:: python
+.. code-block:: cpp
 
     py::type T_py = py::type::of<T>();
 
@@ -1259,3 +1259,40 @@ object, just like ``type(ob)`` in Python.
     Other types, like ``py::type::of<int>()``, do not work, see :ref:`type-conversions`.
 
 .. versionadded:: 2.6
+
+Custom type setup
+=================
+
+For advanced use cases, such as setting certain type object slots, you may wish
+to directly manipulate the `PyHeapTypeObject` corresponding to a ``py::class_``
+definition.
+
+You can do that using ``py::custom_type_setup``:
+
+.. code-block:: cpp
+
+   struct OwnsPythonObjects {
+       py::object value = py::none();
+   };
+   py::class_<OwnsPythonObjects> cls(
+       m, "OwnsPythonObjects", py::custom_type_setup([](PyHeapTypeObject *heap_type) {
+           auto *type = &heap_type->ht_type;
+           type->tp_flags |= Py_TPFLAGS_HAVE_GC;
+           type->tp_traverse = +[](PyObject *self_base, visitproc visit, void *arg) {
+               auto &self = py::cast<OwnsPythonObjects&>(py::handle(self_base));
+               Py_VISIT(self.value.ptr());
+               return 0;
+           };
+           type->tp_clear = +[](PyObject *self_base) {
+               auto &self = py::cast<OwnsPythonObjects&>(py::handle(self_base));
+               // Use temporary object to ensure `self.value` remains valid
+               // while Python APIs are called.
+               py::object temp = py::none();
+               std::swap(temp, self.value);
+               return 0;
+           };
+       }));
+   cls.def(py::init([] { return OwnsPythonObjects{}; }));
+   cls.def_readwrite("value", &OwnsPythonObjects::value);
+
+.. versionadded:: 2.8

include/pybind11/attr.h

@@ -11,6 +11,7 @@
 #pragma once
 
 #include "cast.h"
+#include "detail/function_view.h"
 
 PYBIND11_NAMESPACE_BEGIN(PYBIND11_NAMESPACE)
 
@@ -79,6 +80,23 @@ struct metaclass {
     explicit metaclass(handle value) : value(value) { }
 };
 
+/// Specifies a custom callback with signature `void (PyHeapTypeObject*)` that
+/// may be used to customize the Python type.
+///
+/// The callback is invoked immediately before `PyType_Ready`.
+///
+/// Note: This is an advanced interface, and uses of it may require changes to
+/// work with later versions of pybind11.  You may wish to consult the
+/// implementation of `make_new_python_type` in `detail/classes.h` to understand
+/// the context in which the callback will be run.
+struct custom_type_setup {
+    using callback = detail::function_view<void(PyHeapTypeObject* heap_type)>;
+
+    explicit custom_type_setup(callback value) : value(value) {}
+
+    callback value;
+};
+
 /// Annotation that marks a class as local to the module:
 struct module_local { const bool value;
     constexpr explicit module_local(bool v = true) : value(v) {}
@@ -272,6 +290,9 @@ struct type_record {
     /// Custom metaclass (optional)
     handle metaclass;
 
+    /// Custom type setup.
+    ::pybind11::custom_type_setup::callback custom_type_setup;
+
     /// Multiple inheritance marker
     bool multiple_inheritance : 1;
 
@@ -476,6 +497,11 @@ struct process_attribute<dynamic_attr> : process_attribute_default<dynamic_attr>
     static void init(const dynamic_attr &, type_record *r) { r->dynamic_attr = true; }
 };
 
+template <>
+struct process_attribute<custom_type_setup> {
+    static void init(const custom_type_setup &value, type_record *r) { r->custom_type_setup = value.value; }
+};
+
 template <>
 struct process_attribute<is_final> : process_attribute_default<is_final> {
     static void init(const is_final &, type_record *r) { r->is_final = true; }

include/pybind11/detail/class.h

@@ -683,11 +683,13 @@ inline PyObject* make_new_python_type(const type_record &rec) {
     if (rec.buffer_protocol)
         enable_buffer_protocol(heap_type);
 
+    if (rec.custom_type_setup)
+        rec.custom_type_setup(heap_type);
+
     if (PyType_Ready(type) < 0)
         pybind11_fail(std::string(rec.name) + ": PyType_Ready failed (" + error_string() + ")!");
 
-    assert(rec.dynamic_attr ? PyType_HasFeature(type, Py_TPFLAGS_HAVE_GC)
-                            : !PyType_HasFeature(type, Py_TPFLAGS_HAVE_GC));
+    assert(!rec.dynamic_attr || PyType_HasFeature(type, Py_TPFLAGS_HAVE_GC));
 
     /* Register type with the parent scope */
     if (rec.scope)

include/pybind11/detail/common.h

@@ -570,6 +570,13 @@ template <typename T> using remove_cv_t = typename std::remove_cv<T>::type;
 template <typename T> using remove_reference_t = typename std::remove_reference<T>::type;
 #endif
 
+#ifdef __cpp_lib_remove_cvref
+using std::remove_cvref_t;
+#else
+template <typename T>
+using remove_cvref_t = remove_cv_t<remove_reference_t<T>>;
+#endif
+
 /// Index sequences
 #if defined(PYBIND11_CPP14)
 using std::index_sequence;
@@ -771,6 +778,21 @@ using function_signature_t = conditional_t<
 template <typename T> using is_lambda = satisfies_none_of<remove_reference_t<T>,
         std::is_function, std::is_pointer, std::is_member_pointer>;
 
+template <typename F, typename... Arg>
+using call_result_t = decltype(std::declval<F>()(std::declval<Arg>()...));
+
+template <typename Expected>
+struct is_expected_return_type {
+    template <typename Actual>
+    using compatible_with = std::is_convertible<Actual, Expected>;
+};
+
+template <>
+struct is_expected_return_type<void> {
+    template <typename Actual>
+    using compatible_with = std::true_type;
+};
+
 // [workaround(intel)] Internal error on fold expression
 /// Apply a function over each element of a parameter pack
 #if defined(__cpp_fold_expressions) && !defined(__INTEL_COMPILER)

include/pybind11/detail/function_view.h

@@ -0,0 +1,59 @@
+// Copyright (c) 2021 The Pybind Development Team.
+// 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 "common.h"
+
+PYBIND11_NAMESPACE_BEGIN(PYBIND11_NAMESPACE)
+PYBIND11_NAMESPACE_BEGIN(detail)
+
+template <typename Signature>
+class function_view;
+
+template <typename R, typename... Arg>
+class function_view<R(Arg...)> {
+public:
+    function_view() = default;
+
+    /// Creates a view that references an existing function.
+    ///
+    /// The constructed view is only valid for the lifetime of `f`.
+    ///
+    /// Requires that `F` is a function object type that is callable with
+    /// arguments `(Args...)` and has a return type convertible to `R`.
+    template <typename F,
+              typename = enable_if_t<
+                  // Prevent construction from the same `function_view` type, to avoid inefficient
+                  // double wrapping that also commonly leads to dangling references.
+                  (!std::is_same<remove_cvref_t<F>, function_view>::value
+                   // Ensure `F` is compatible with the signature.
+                   && is_expected_return_type<R>::template compatible_with<
+                       call_result_t<F, Arg...>>::value)>>
+    function_view(F &&f) noexcept // NOLINT(google-explicit-constructor)
+        : erased_fn_(&wrapper<remove_reference_t<F>>), erased_obj_(&f) {}
+
+    /// Calls the referenced function.
+    ///
+    /// Precondition: this is a non-null function view.
+    R operator()(Arg... arg) const {
+        return erased_fn_(const_cast<void *>(erased_obj_), static_cast<Arg &&>(arg)...);
+    }
+
+    /// Returns `true` if this is a non-null function view.
+    explicit operator bool() const { return erased_fn_ != nullptr; }
+
+private:
+    template <typename F>
+    static R wrapper(void *obj, Arg... arg) {
+        return (*static_cast<typename std::add_pointer<F>::type>(obj))(
+            static_cast<Arg &&>(arg)...);
+    }
+
+    R (*erased_fn_)(void *, Arg...) = nullptr;
+    const void *erased_obj_;
+};
+
+PYBIND11_NAMESPACE_END(detail)
+PYBIND11_NAMESPACE_END(PYBIND11_NAMESPACE)

tests/CMakeLists.txt

@@ -104,6 +104,7 @@ set(PYBIND11_TEST_FILES
     test_constants_and_functions.cpp
     test_copy_move.cpp
     test_custom_type_casters.cpp
+    test_custom_type_setup.cpp
     test_docstring_options.cpp
     test_eigen.cpp
     test_enum.cpp

tests/extra_python_package/test_files.py

@@ -40,6 +40,7 @@
     "include/pybind11/detail/class.h",
     "include/pybind11/detail/common.h",
     "include/pybind11/detail/descr.h",
+    "include/pybind11/detail/function_view.h",
     "include/pybind11/detail/init.h",
     "include/pybind11/detail/internals.h",
     "include/pybind11/detail/type_caster_base.h",

tests/test_custom_type_setup.cpp

@@ -0,0 +1,41 @@
+/*
+    tests/test_custom_type_setup.cpp -- Tests `pybind11::custom_type_setup`
+
+    Copyright (c) Google LLC
+
+    All rights reserved. Use of this source code is governed by a
+    BSD-style license that can be found in the LICENSE file.
+*/
+
+#include <pybind11/pybind11.h>
+
+#include "pybind11_tests.h"
+
+namespace py = pybind11;
+
+namespace {
+
+struct OwnsPythonObjects {
+    py::object value = py::none();
+};
+} // namespace
+
+TEST_SUBMODULE(custom_type_setup, m) {
+    py::class_<OwnsPythonObjects> cls(
+        m, "OwnsPythonObjects", py::custom_type_setup([](PyHeapTypeObject *heap_type) {
+            auto *type = &heap_type->ht_type;
+            type->tp_flags |= Py_TPFLAGS_HAVE_GC;
+            type->tp_traverse = [](PyObject *self_base, visitproc visit, void *arg) {
+                auto &self = py::cast<OwnsPythonObjects &>(py::handle(self_base));
+                Py_VISIT(self.value.ptr());
+                return 0;
+            };
+            type->tp_clear = [](PyObject *self_base) {
+                auto &self = py::cast<OwnsPythonObjects &>(py::handle(self_base));
+                self.value = py::none();
+                return 0;
+            };
+        }));
+    cls.def(py::init([] { return OwnsPythonObjects{}; }));
+    cls.def_readwrite("value", &OwnsPythonObjects::value);
+}

tests/test_custom_type_setup.py

@@ -0,0 +1,50 @@
+# -*- coding: utf-8 -*-
+
+import gc
+import weakref
+
+import pytest
+
+import env  # noqa: F401
+from pybind11_tests import custom_type_setup as m
+
+
+@pytest.fixture
+def gc_tester():
+    """Tests that an object is garbage collected."""
+
+    weak_refs = []
+
+    def add_ref(obj):
+        # PyPy does not support `gc.is_tracked`.
+        if hasattr(gc, "is_tracked"):
+            assert gc.is_tracked(obj)
+        weak_refs.append(weakref.ref(obj))
+
+    yield add_ref
+
+    gc.collect()
+    if hasattr(gc, "collect_step"):
+        # PyPy may require additional encouragement to actually collect the
+        # garbage.
+        for _ in range(100):
+            gc.collect_step()
+    for ref in weak_refs:
+        assert ref() is None
+
+
+# PyPy does not seem to garbage collect.
+@pytest.mark.xfail("env.PYPY")
+def test_self_cycle(gc_tester):
+    obj = m.OwnsPythonObjects()
+    obj.value = obj
+    gc_tester(obj)
+
+
+# PyPy does not seem to garbage collect.
+@pytest.mark.xfail("env.PYPY")
+def test_indirect_cycle(gc_tester):
+    obj = m.OwnsPythonObjects()
+    obj_list = [obj]
+    obj.value = obj_list
+    gc_tester(obj)