From af936e1987312e665b64463bed94d4a486a78a09 Mon Sep 17 00:00:00 2001 From: Matthieu Bec Date: Fri, 3 Mar 2017 08:45:50 -0800 Subject: [PATCH] Expose enum_ entries as "__members__" read-only property. Getters get a copy. --- docs/classes.rst | 6 ++++++ include/pybind11/pybind11.h | 39 ++++++++++++++++++++++--------------- tests/test_enum.py | 9 +++++++++ 3 files changed, 38 insertions(+), 16 deletions(-) diff --git a/docs/classes.rst b/docs/classes.rst index 2c1ff2a39..30fb2a2d5 100644 --- a/docs/classes.rst +++ b/docs/classes.rst @@ -423,6 +423,12 @@ typed enums. >>> int(p.type) 1L +The entries defined by the enumeration type are exposed in the ``__members__`` property: + +.. code-block:: pycon + + >>> Pet.Kind.__members__ + {'Dog': Kind.Dog, 'Cat': Kind.Cat} .. note:: diff --git a/include/pybind11/pybind11.h b/include/pybind11/pybind11.h index 242ae2af3..580c4ef9f 100644 --- a/include/pybind11/pybind11.h +++ b/include/pybind11/pybind11.h @@ -1116,24 +1116,32 @@ private: template class enum_ : public class_ { public: using class_::def; + using class_::def_property_readonly_static; using Scalar = typename std::underlying_type::type; template using arithmetic_tag = std::is_same; template enum_(const handle &scope, const char *name, const Extra&... extra) - : class_(scope, name, extra...), m_parent(scope) { + : class_(scope, name, extra...), m_entries(), m_parent(scope) { constexpr bool is_arithmetic = !std::is_same, void>::value; - auto entries = new std::unordered_map(); - def("__repr__", [name, entries](Type value) -> std::string { - auto it = entries->find((Scalar) value); - return std::string(name) + "." + - ((it == entries->end()) ? std::string("???") - : std::string(it->second)); + auto m_entries_ptr = m_entries.inc_ref().ptr(); + def("__repr__", [name, m_entries_ptr](Type value) -> pybind11::str { + for (const auto &kv : reinterpret_borrow(m_entries_ptr)) { + if (pybind11::cast(kv.second) == value) + return pybind11::str("{}.{}").format(name, kv.first); + } + return pybind11::str("{}.???").format(name); }); + def_property_readonly_static("__members__", [m_entries_ptr](object /* self */) { + dict m; + for (const auto &kv : reinterpret_borrow(m_entries_ptr)) + m[kv.first] = kv.second; + return m; + }, return_value_policy::copy); def("__init__", [](Type& value, Scalar i) { value = (Type)i; }); def("__init__", [](Type& value, Scalar i) { new (&value) Type((Type) i); }); def("__int__", [](Type value) { return (Scalar) value; }); @@ -1172,26 +1180,25 @@ public: // Pickling and unpickling -- needed for use with the 'multiprocessing' module def("__getstate__", [](const Type &value) { return pybind11::make_tuple((Scalar) value); }); def("__setstate__", [](Type &p, tuple t) { new (&p) Type((Type) t[0].cast()); }); - m_entries = entries; } /// Export enumeration entries into the parent scope - enum_ &export_values() { - for (auto item : reinterpret_borrow(((PyTypeObject *) this->m_ptr)->tp_dict)) { - if (isinstance(item.second, this->m_ptr)) - m_parent.attr(item.first) = item.second; - } + enum_& export_values() { + for (const auto &kv : m_entries) + m_parent.attr(kv.first) = kv.second; return *this; } /// Add an enumeration entry enum_& value(char const* name, Type value) { - this->attr(name) = pybind11::cast(value, return_value_policy::copy); - (*m_entries)[(Scalar) value] = name; + auto v = pybind11::cast(value, return_value_policy::copy); + this->attr(name) = v; + m_entries[pybind11::str(name)] = v; return *this; } + private: - std::unordered_map *m_entries; + dict m_entries; handle m_parent; }; diff --git a/tests/test_enum.py b/tests/test_enum.py index de5f3c6f6..ba7e22ab0 100644 --- a/tests/test_enum.py +++ b/tests/test_enum.py @@ -7,6 +7,15 @@ def test_unscoped_enum(): assert str(UnscopedEnum.EOne) == "UnscopedEnum.EOne" assert str(UnscopedEnum.ETwo) == "UnscopedEnum.ETwo" assert str(EOne) == "UnscopedEnum.EOne" + # __members__ property + assert UnscopedEnum.__members__ == {"EOne": UnscopedEnum.EOne, "ETwo": UnscopedEnum.ETwo} + # __members__ readonly + with pytest.raises(AttributeError): + UnscopedEnum.__members__ = {} + # __members__ returns a copy + foo = UnscopedEnum.__members__ + foo["bar"] = "baz" + assert UnscopedEnum.__members__ == {"EOne": UnscopedEnum.EOne, "ETwo": UnscopedEnum.ETwo} # no TypeError exception for unscoped enum ==/!= int comparisons y = UnscopedEnum.ETwo