1 files changed, 155 insertions, 0 deletions
diff --git a/libs/plugin_python/src/pybind11-utils/include/pybind11_utils/functional.h b/libs/plugin_python/src/pybind11-utils/include/pybind11_utils/functional.h
new file mode 100644
index 0000000..8662d82
--- /dev/null
+++ b/libs/plugin_python/src/pybind11-utils/include/pybind11_utils/functional.h
@@ -0,0 +1,155 @@
+#ifndef PYTHON_PYBIND11_FUNCTIONAL_H
+#define PYTHON_PYBIND11_FUNCTIONAL_H
+
+#include <pybind11/pybind11.h>
+
+namespace mo2::python::detail {
+
+    // check if the given function is valid for a C++ function with the
+    // given arity
+    //
+    bool has_compatible_arity(pybind11::function handle, std::size_t arity);
+
+}  // namespace mo2::python::detail
+
+namespace pybind11::detail {
+
+    // custom type_caster for std::function<>
+    //
+    // most of this is from pybind11 except that we also check arity of the function to
+    // allow overloaded function based on arity of argument
+    //
+    template <typename Return, typename... Args>
+    struct type_caster<std::function<Return(Args...)>> {
+        using type = std::function<Return(Args...)>;
+        using retval_type =
+            conditional_t<std::is_same<Return, void>::value, void_type, Return>;
+        using function_type = Return (*)(Args...);
+
+    public:
+        bool load(handle src, bool convert)
+        {
+            if (src.is_none()) {
+                // Defer accepting None to other overloads (if we aren't in convert
+                // mode):
+                if (!convert) {
+                    return false;
+                }
+                return true;
+            }
+
+            if (!isinstance<function>(src)) {
+                return false;
+            }
+
+            auto func = reinterpret_borrow<function>(src);
+
+            /*
+               When passing a C++ function as an argument to another C++
+               function via Python, every function call would normally involve
+               a full C++ -> Python -> C++ roundtrip, which can be prohibitive.
+               Here, we try to at least detect the case where the function is
+               stateless (i.e. function pointer or lambda function without
+               captured variables), in which case the roundtrip can be avoided.
+             */
+            if (auto cfunc = func.cpp_function()) {
+                auto* cfunc_self = PyCFunction_GET_SELF(cfunc.ptr());
+                if (isinstance<capsule>(cfunc_self)) {
+                    auto c    = reinterpret_borrow<capsule>(cfunc_self);
+                    auto* rec = (function_record*)c;
+
+                    while (rec != nullptr) {
+                        if (rec->is_stateless &&
+                            same_type(typeid(function_type),
+                                      *reinterpret_cast<const std::type_info*>(
+                                          rec->data[1]))) {
+                            struct capture {
+                                function_type f;
+                            };
+                            value = ((capture*)&rec->data)->f;
+                            return true;
+                        }
+                        rec = rec->next;
+                    }
+                }
+                // PYPY segfaults here when passing builtin function like sum.
+                // Raising an fail exception here works to prevent the segfault, but
+                // only on gcc. See PR #1413 for full details
+            }
+
+            // !MO2! - check arity
+
+            if (!mo2::python::detail::has_compatible_arity(func, sizeof...(Args))) {
+                return false;
+            }
+
+            // !MO2! - everything below is copy/paste from pybind11
+
+            // ensure GIL is held during functor destruction
+            struct func_handle {
+                function f;
+#if !(defined(_MSC_VER) && _MSC_VER == 1916 && defined(PYBIND11_CPP17))
+                // This triggers a syntax error under very special conditions (very
+                // weird indeed).
+                explicit
+#endif
+                    func_handle(function&& f_) noexcept
+                    : f(std::move(f_))
+                {
+                }
+                func_handle(const func_handle& f_) { operator=(f_); }
+                func_handle& operator=(const func_handle& f_)
+                {
+                    gil_scoped_acquire acq;
+                    f = f_.f;
+                    return *this;
+                }
+                ~func_handle()
+                {
+                    gil_scoped_acquire acq;
+                    function kill_f(std::move(f));
+                }
+            };
+
+            // to emulate 'move initialization capture' in C++11
+            struct func_wrapper {
+                func_handle hfunc;
+                explicit func_wrapper(func_handle&& hf) noexcept : hfunc(std::move(hf))
+                {
+                }
+                Return operator()(Args... args) const
+                {
+                    gil_scoped_acquire acq;
+                    object retval(hfunc.f(std::forward<Args>(args)...));
+                    return retval.template cast<Return>();
+                }
+            };
+
+            value = func_wrapper(func_handle(std::move(func)));
+            return true;
+        }
+
+        template <typename Func>
+        static handle cast(Func&& f_, return_value_policy policy, handle /* parent */)
+        {
+            if (!f_) {
+                return none().inc_ref();
+            }
+
+            auto result = f_.template target<function_type>();
+            if (result) {
+                return cpp_function(*result, policy).release();
+            }
+            return cpp_function(std::forward<Func>(f_), policy).release();
+        }
+
+        PYBIND11_TYPE_CASTER(type, const_name("Callable[[") +
+                                       concat(make_caster<Args>::name...) +
+                                       const_name("], ") +
+                                       make_caster<retval_type>::name +
+                                       const_name("]"));
+    };
+
+}  // namespace pybind11::detail
+
+#endif