diff --git a/dtool/src/interrogatedb/py_panda.cxx b/dtool/src/interrogatedb/py_panda.cxx index 7c2bef1ce7..6ebfff5b4b 100644 --- a/dtool/src/interrogatedb/py_panda.cxx +++ b/dtool/src/interrogatedb/py_panda.cxx @@ -286,6 +286,43 @@ PyObject *_Dtool_Return(PyObject *value) { } #if PY_VERSION_HEX < 0x03040000 +/** + * This function converts an int value to the appropriate enum instance. + */ +PyObject *Dtool_EnumType_New(PyTypeObject *subtype, PyObject *args, PyObject *kwds) { + PyObject *arg; + if (!Dtool_ExtractArg(&arg, args, kwds, "value")) { + return PyErr_Format(PyExc_TypeError, + "%s() missing 1 required argument: 'value'", + subtype->tp_name); + } + + if (Py_TYPE(arg) == subtype) { + Py_INCREF(arg); + return arg; + } + + PyObject *value2member = PyDict_GetItemString(subtype->tp_dict, "_value2member_map_"); + nassertr_always(value2member != nullptr, nullptr); + + PyObject *member = PyDict_GetItem(value2member, arg); + if (member != nullptr) { + Py_INCREF(member); + return member; + } + + PyObject *repr = PyObject_Repr(arg); + PyErr_Format(PyExc_ValueError, "%s is not a valid %s", +#if PY_MAJOR_VERSION >= 3 + PyUnicode_AS_STRING(repr), +#else + PyString_AS_STRING(repr), +#endif + subtype->tp_name); + Py_DECREF(repr); + return nullptr; +} + static PyObject *Dtool_EnumType_Str(PyObject *self) { PyObject *name = PyObject_GetAttrString(self, "name"); #if PY_MAJOR_VERSION >= 3 @@ -337,6 +374,7 @@ PyTypeObject *Dtool_EnumType_Create(const char *name, PyObject *names, const cha static PyObject *name_sunder_str; static PyObject *value_str; static PyObject *value_sunder_str; + static PyObject *value2member_map_sunder_str; // Emulate something vaguely like the enum module. if (enum_class == nullptr) { #if PY_MAJOR_VERSION >= 3 @@ -344,11 +382,13 @@ PyTypeObject *Dtool_EnumType_Create(const char *name, PyObject *names, const cha value_str = PyUnicode_InternFromString("value"); name_sunder_str = PyUnicode_InternFromString("_name_"); value_sunder_str = PyUnicode_InternFromString("_value_"); + value2member_map_sunder_str = PyUnicode_InternFromString("_value2member_map_"); #else name_str = PyString_InternFromString("name"); value_str = PyString_InternFromString("value"); name_sunder_str = PyString_InternFromString("_name_"); value_sunder_str = PyString_InternFromString("_value_"); + value2member_map_sunder_str = PyString_InternFromString("_value2member_map_"); #endif PyObject *name_value_tuple = PyTuple_New(4); PyTuple_SET_ITEM(name_value_tuple, 0, name_str); @@ -365,27 +405,39 @@ PyTypeObject *Dtool_EnumType_Create(const char *name, PyObject *names, const cha enum_class = PyObject_CallFunction((PyObject *)&PyType_Type, (char *)"s()N", "Enum", slots_dict); nassertr(enum_class != nullptr, nullptr); } - PyObject *result = PyObject_CallFunction((PyObject *)&PyType_Type, (char *)"s(O)N", name, enum_class, PyDict_New()); + + // Create a subclass of this generic Enum class we just created. + PyObject *value2member = PyDict_New(); + PyObject *dict = PyDict_New(); + PyDict_SetItem(dict, value2member_map_sunder_str, value2member); + PyObject *result = PyObject_CallFunction((PyObject *)&PyType_Type, (char *)"s(O)N", name, enum_class, dict); nassertr(result != nullptr, nullptr); + ((PyTypeObject *)result)->tp_new = Dtool_EnumType_New; ((PyTypeObject *)result)->tp_str = Dtool_EnumType_Str; ((PyTypeObject *)result)->tp_repr = Dtool_EnumType_Repr; - // Copy the names as instances of the above to the class dict. + PyObject *empty_tuple = PyTuple_New(0); + + // Copy the names as instances of the above to the class dict, and create a + // reverse mapping in the _value2member_map_ dict. Py_ssize_t size = PyTuple_GET_SIZE(names); for (Py_ssize_t i = 0; i < size; ++i) { PyObject *item = PyTuple_GET_ITEM(names, i); PyObject *name = PyTuple_GET_ITEM(item, 0); PyObject *value = PyTuple_GET_ITEM(item, 1); - PyObject *member = _PyObject_CallNoArg(result); + PyObject *member = PyType_GenericNew((PyTypeObject *)result, empty_tuple, nullptr); PyObject_SetAttr(member, name_str, name); PyObject_SetAttr(member, name_sunder_str, name); PyObject_SetAttr(member, value_str, value); PyObject_SetAttr(member, value_sunder_str, value); PyObject_SetAttr(result, name, member); + PyDict_SetItem(value2member, value, member); Py_DECREF(member); } Py_DECREF(names); + Py_DECREF(value2member); + Py_DECREF(empty_tuple); #endif if (module != nullptr) {