Implement py3_enum<> mapping to enum.IntEnum

Author: aldanor

include/pybind11/cast.h

@@ -890,6 +890,77 @@ template <typename... Tuple> class type_caster<std::tuple<Tuple...>> {
     std::tuple<make_caster<Tuple>...> value;
 };
 
+struct py3_enum_info {
+    handle type = {};
+    std::unordered_map<long long, handle> values = {};
+
+    py3_enum_info() = default;
+
+    py3_enum_info(handle type, const dict& values) : type(type) {
+        for (auto item : values)
+            this->values[static_cast<long long>(item.second.cast<int>())] = type.attr(item.first);
+    }
+
+    static std::unordered_map<std::type_index, py3_enum_info>& registry() {
+        static std::unordered_map<std::type_index, py3_enum_info> map = {};
+        return map;
+    }
+
+    template<typename T>
+    static void bind(handle type, const dict& values) {
+        registry()[typeid(T)] = py3_enum_info(type, values);
+    }
+
+    template<typename T>
+    static const py3_enum_info* get() {
+        auto it = registry().find(typeid(T));
+        return it == registry().end() ? nullptr : &it->second;
+    }
+};
+
+template<typename T>
+struct type_caster<T, enable_if_t<std::is_enum<T>::value>> {
+private:
+    using base_caster = type_caster_base<T>;
+    base_caster caster;
+    bool py3 = false;
+    T value;
+
+public:
+    template<typename U> using cast_op_type = pybind11::detail::cast_op_type<U>;
+
+    operator T*() { return py3 ? &value : static_cast<T*>(caster); }
+    operator T&() { return py3 ? value : static_cast<T&>(caster); }
+
+    static handle cast(const T& src, return_value_policy rvp, handle parent) {
+        if (auto info = py3_enum_info::get<T>()) {
+            auto it = info->values.find(static_cast<long long>(src));
+            if (it == info->values.end())
+                return {};
+            return it->second.inc_ref();
+        }
+        return base_caster::cast(src, rvp, parent);
+    }
+
+    bool load(handle src, bool convert) {
+        if (!src)
+            return false;
+        if (auto info = py3_enum_info::get<T>()) {
+            py3 = true;
+            if (!isinstance(src, info->type))
+                return false;
+            value = static_cast<T>(src.cast<long long>());
+            return true;
+        }
+        py3 = false;
+        return caster.load(src, convert);
+    }
+
+    static PYBIND11_DESCR name() {
+        return base_caster::name();
+    }
+};
+
 /// Helper class which abstracts away certain actions. Users can provide specializations for
 /// custom holders, but it's only necessary if the type has a non-standard interface.
 template <typename T>

include/pybind11/pybind11.h

@@ -1238,6 +1238,39 @@ template <typename Type> class enum_ : public class_<Type> {
     handle m_parent;
 };
 
+template<typename T>
+class py3_enum {
+public:
+    using underlying_type = typename std::underlying_type<T>::type;
+
+    py3_enum(handle scope, const char* name)
+    : name(name),
+      parent(scope),
+      ctor(module::import("enum").attr("IntEnum")),
+      unique(module::import("enum").attr("unique")) {
+        update();
+    }
+
+    py3_enum& value(const char* name, T value) {
+        entries[name] = cast(static_cast<underlying_type>(value));
+        update();
+        return *this;
+    }
+
+private:
+    const char *name;
+    handle parent;
+    dict entries;
+    object ctor;
+    object unique;
+
+    void update() {
+        object type = unique(ctor(name, entries));
+        setattr(parent, name, type);
+        detail::py3_enum_info::bind<T>(type, entries);
+    }
+};
+
 NAMESPACE_BEGIN(detail)
 template <typename... Args> struct init {
     template <typename Class, typename... Extra, enable_if_t<!Class::has_alias, int> = 0>