Add custom_type_setup attribute (#3287)

* Fix `pybind11::object::operator=` to be safe if `*this` is accessible from Python

* 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.

docs/advanced/classes.rst

@@ -1261,3 +1261,37 @@ 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 enabling garbage collection support, 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));
+               self.value = py::none();
+               return 0;
+           };
+       }));
+   cls.def(py::init<>());
+   cls.def_readwrite("value", &OwnsPythonObjects::value);
+
+.. versionadded:: 2.8

include/pybind11/attr.h

@@ -12,6 +12,8 @@
 
 #include "cast.h"
 
+#include <functional>
+
 PYBIND11_NAMESPACE_BEGIN(PYBIND11_NAMESPACE)
 
 /// \addtogroup annotations
@@ -79,6 +81,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 = std::function<void(PyHeapTypeObject *heap_type)>;
+
+    explicit custom_type_setup(callback value) : value(std::move(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 +291,9 @@ struct type_record {
     /// Custom metaclass (optional)
     handle metaclass;
 
+    /// Custom type setup.
+    custom_type_setup::callback custom_type_setup_callback;
+
     /// Multiple inheritance marker
     bool multiple_inheritance : 1;
 
@@ -476,6 +498,13 @@ 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_callback = 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_callback)
+        rec.custom_type_setup_callback(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/pytypes.h

@@ -259,8 +259,11 @@ public:
 
     object& operator=(const object &other) {
         other.inc_ref();
-        dec_ref();
+        // Use temporary variable to ensure `*this` remains valid while
+        // `Py_XDECREF` executes, in case `*this` is accessible from Python.
+        handle temp(m_ptr);
         m_ptr = other.m_ptr;
+        temp.dec_ref();
         return *this;
     }
 

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/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<>());
+    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.
+
+    Assumes that any unreferenced objects are fully collected after calling
+    `gc.collect()`.  That is true on CPython, but does not appear to reliably
+    hold on PyPy.
+    """
+
+    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()
+    for ref in weak_refs:
+        assert ref() is None
+
+
+# PyPy does not seem to reliably garbage collect.
+@pytest.mark.skipif("env.PYPY")
+def test_self_cycle(gc_tester):
+    obj = m.OwnsPythonObjects()
+    obj.value = obj
+    gc_tester(obj)
+
+
+# PyPy does not seem to reliably garbage collect.
+@pytest.mark.skipif("env.PYPY")
+def test_indirect_cycle(gc_tester):
+    obj = m.OwnsPythonObjects()
+    obj_list = [obj]
+    obj.value = obj_list
+    gc_tester(obj)