Merge pull request #187 from quantopian/autoclass-fixes

ENH: Autoclass fixes

Author: gerrymanoim

include/libpy/autoclass.h

@@ -89,6 +89,7 @@ class autoclass_impl {
     py::owned_ref<PyTypeObject> m_type;
     PyType_Spec m_spec;
     py::owned_ref<PyTypeObject> m_py_basetype;
+    py::owned_ref<> m_module;
 
     /** Check if this type uses the `Py_TPFLAGS_HAVE_GC`, which requires that we implement
         at least `Py_tp_traverse`, and will use `PyObject_GC_New` and `PyObject_GC_Del`.
@@ -153,32 +154,32 @@ class autoclass_impl {
 
     // dispatch for free function that accepts as a first argument `T`
     template<typename R, typename... Args, auto impl>
-    struct free_function_impl<R(*)(T, Args...), impl>
+    struct free_function_impl<R (*)(T, Args...), impl>
         : public free_function_base<impl, R, Args...> {};
 
     // dispatch for free function that accepts as a first argument `T&`
     template<typename R, typename... Args, auto impl>
-    struct free_function_impl<R(*)(T&, Args...), impl>
+    struct free_function_impl<R (*)(T&, Args...), impl>
         : public free_function_base<impl, R, Args...> {};
 
     // dispatch for free function that accepts as a first argument `const T&`
     template<typename R, typename... Args, auto impl>
-    struct free_function_impl<R(*)(const T&, Args...), impl>
+    struct free_function_impl<R (*)(const T&, Args...), impl>
         : public free_function_base<impl, R, Args...> {};
 
     // dispatch for a noexcept free function that accepts as a first argument `T`
     template<typename R, typename... Args, auto impl>
-    struct free_function_impl<R(*)(T, Args...) noexcept, impl>
+    struct free_function_impl<R (*)(T, Args...) noexcept, impl>
         : public free_function_base<impl, R, Args...> {};
 
     // dispatch for noexcept free function that accepts as a first argument `T&`
     template<typename R, typename... Args, auto impl>
-    struct free_function_impl<R(*)(T&, Args...) noexcept, impl>
+    struct free_function_impl<R (*)(T&, Args...) noexcept, impl>
         : public free_function_base<impl, R, Args...> {};
 
     // dispatch for a noexcept free function that accepts as a first argument `const T&`
     template<typename R, typename... Args, auto impl>
-    struct free_function_impl<R(*)(const T&, Args...) noexcept, impl>
+    struct free_function_impl<R (*)(const T&, Args...) noexcept, impl>
         : public free_function_base<impl, R, Args...> {};
 
     template<auto impl, typename R, typename... Args>
@@ -336,19 +337,42 @@ class autoclass_impl {
         return static_cast<void*>(std::addressof(unbox(ob)));
     }
 
+private:
+    std::string name_in_module(py::borrowed_ref<> module, std::string_view name) {
+        if (!module) {
+            return std::string{name};
+        }
+
+        const char* const module_name = PyModule_GetName(module.get());
+        if (!module_name) {
+            throw py::exception{};
+        }
+        return py::util::format_string(module_name, ".", name);
+    }
+
 public:
-    autoclass_impl(std::string name = util::type_name<T>(),
+    /** Construct the type and add it to a module.
+
+        @param module The module to add the type to.
+        @param name The name of the type as seen from Python.
+        @param extra_flags Extra flags to forward to `tp_flags` field.
+        @param base_type A Python type to subclass.
+     */
+    autoclass_impl(py::borrowed_ref<> module,
+                   std::string name = py::util::type_name<T>(),
                    int extra_flags = 0,
                    py::borrowed_ref<PyTypeObject> base_type = nullptr)
-        : m_storage(std::make_unique<detail::autoclass_storage>(dynamic_unbox,
-                                                                std::move(name))),
+        : m_storage(
+              std::make_unique<detail::autoclass_storage>(dynamic_unbox,
+                                                          name_in_module(module, name))),
           m_type(nullptr),
           m_spec({m_storage->strings.front().data(),
                   static_cast<int>(sizeof(object)),
                   0,
                   flags(extra_flags, base_type),
                   nullptr}),
-          m_py_basetype(py::owned_ref<PyTypeObject>::xnew_reference(base_type)) {
+          m_py_basetype(py::owned_ref<PyTypeObject>::xnew_reference(base_type)),
+          m_module(py::owned_ref<>::xnew_reference(module)) {
         if (base_type) {
             // Check to make sure that the static base type is not obviously
             // wrong. This check does not ensure that the static base type is
@@ -402,18 +426,23 @@ class autoclass_impl {
         add_slot(Py_tp_dealloc, py_dealloc);
     }
 
-    // Delete the copy constructor, the intermediate string data points into
-    // storage that is managed by the type until `.type()` is called.
-    // Also, don't try to create 2 versions of the same type.
+    autoclass_impl(std::string name = util::type_name<T>(),
+                   int extra_flags = 0,
+                   py::borrowed_ref<PyTypeObject> base_type = nullptr)
+        : autoclass_impl(nullptr, std::move(name), extra_flags, base_type) {}
+
+    // Delete the copy constructor, the intermediate string data points into storage
+    // that is managed by the type until `.type()` is called. Also, don't try to
+    // create 2 versions of the same type.
     autoclass_impl(const autoclass_impl&) = delete;
     autoclass_impl(autoclass_impl&&) = default;
     autoclass_impl& operator=(autoclass_impl&&) = default;
 
     /** Add a `tp_traverse` field to this type. This is only allowed, but required if
         `extra_flags & Py_TPFLAGS_HAVE_GC`.
 
-        @tparam impl The implementation of the traverse function. This should either be an
-                     `int(T&, visitproc, void*)` or `int (T::*)(visitproc, void*)`.
+        @tparam impl The implementation of the traverse function. This should either
+       be an `int(T&, visitproc, void*)` or `int (T::*)(visitproc, void*)`.
      */
     template<auto impl>
     concrete& traverse() {
@@ -441,8 +470,8 @@ class autoclass_impl {
     /** Add a `tp_clear` field to this type. This is only allowed if
         `extra_flags & Py_TPFLAGS_HAVE_GC`.
 
-        @tparam impl The implementation of the clear function. This should either be an
-                     `int(T&)` or `int (T::*)()`.
+        @tparam impl The implementation of the clear function. This should either be
+       an `int(T&)` or `int (T::*)()`.
      */
     template<auto impl>
     concrete& clear() {
@@ -1157,13 +1186,12 @@ class autoclass_impl {
         iter_name += "::iterator";
 
         // create the iterator class and put it in the cache
-        if (!autoclass<iter>(std::move(iter_name), Py_TPFLAGS_HAVE_GC)
-                 .add_slot(Py_tp_iternext, static_cast<iternextfunc>(iternext))
-                 .add_slot(Py_tp_iter, &PyObject_SelfIter)
-                 .template traverse<&iter::traverse>()
-                 .type()) {
-            throw py::exception{};
-        }
+        autoclass<iter>(std::move(iter_name), Py_TPFLAGS_HAVE_GC)
+            .add_slot(Py_tp_iternext, static_cast<iternextfunc>(iternext))
+            .add_slot(Py_tp_iter, &PyObject_SelfIter)
+            .template traverse<&iter::traverse>()
+            .type()
+            .escape();
 
         return [](PyObject* self) -> PyObject* {
             try {
@@ -1438,6 +1466,16 @@ class autoclass_impl {
         release_type_cache.dismiss();
 
         m_type = type;
+
+        if (m_module) {
+            const char* const last_dot = std::strrchr(m_type.get()->tp_name, '.');
+            if (!last_dot) {
+                throw py::exception(PyExc_RuntimeError, "no '.' in type name");
+            }
+            PyObject_SetAttrString(m_module.get(),
+                                   last_dot + 1,
+                                   static_cast<PyObject*>(type));
+        }
         return type;
     }
 

tests/_test_automodule.cc

@@ -25,11 +25,11 @@ LIBPY_AUTOMODULE(tests,
                  ({py::autofunction<is_42>("is_42"),
                    py::autofunction<is_true>("is_true")}))
 (py::borrowed_ref<> m) {
-    py::owned_ref t = py::autoclass<int_float_pair>("_test_automodule.int_float_pair")
-                          .new_<int, float>()
-                          .comparisons<int_float_pair>()
-                          .def<first>("first")
-                          .def<second>("second")
-                          .type();
-    return PyObject_SetAttrString(m.get(), "int_float_pair", static_cast<PyObject*>(t));
+    py::autoclass<int_float_pair>(m, "int_float_pair")
+        .new_<int, float>()
+        .comparisons<int_float_pair>()
+        .def<first>("first")
+        .def<second>("second")
+        .type();
+    return false;
 }

tests/test_autoclass.cc

@@ -16,8 +16,10 @@ namespace test_autoclass {
 using namespace std::literals;
 
 class autoclass : public with_python_interpreter {
+protected:
     std::size_t m_cache_start_size;
 
+private:
     void SetUp() override {
         // Ensure no types are hanging out before we check the cache size.
         gc_collect();
@@ -969,6 +971,9 @@ TEST_F(autoclass, iter) {
         int unboxed = py::from_object<int>(PySequence_Fast_GET_ITEM(fast_seq.get(), ix));
         EXPECT_EQ(unboxed, ix);
     }
+
+    // HACK: `iter()` currently leaks the type
+    ++m_cache_start_size;
 }
 
 TEST_F(autoclass, iter_throws) {
@@ -1010,6 +1015,9 @@ TEST_F(autoclass, iter_throws) {
     EXPECT_FALSE(fast_seq);
     expect_pyerr_type_and_message(PyExc_RuntimeError, "a C++ exception was raised: ayy");
     PyErr_Clear();
+
+    // HACK: `iter()` currently leaks the type
+    ++m_cache_start_size;
 }
 #endif  // LIBPY_AUTOCLASS_UNSAFE_API