Add support for Flag enums

docs/api_core.rst

@@ -1926,6 +1926,14 @@ The following annotations can be specified using the variable-length
    mixed enum types (such as ``Shape.Circle + Color.Red``) are
    permissible.
 
+   .. cpp:struct:: is_flag_enum
+
+   Indicate that the enumeration may be used with bitwise
+   operations.  This enables the bitwise operators ``| & ^ ~``
+   with two enumeration as operands.
+   The result will an enumeration of the same type.
+   So ``Shape(2) | Shape(1) -> Shepe(3)`.
+
 Function binding
 ----------------
 

docs/changelog.rst

@@ -15,6 +15,11 @@ case, both modules must use the same nanobind ABI version, or they will be
 isolated from each other. Releases that don't explicitly mention an ABI version
 below inherit that of the preceding release.
 
+Version 2.1.0 (TBA)
+-------------------
+* The :cpp:class:`nb::enum_\<T\>() <enum_>` can now create an``enum.Flag``-derived type
+  with  flag :cpp:class:`nb::is_flag_enum() <is_flag_enum>`.
+
 Version 2.0.1 (TBA)
 ---------------------------
 
@@ -138,9 +143,11 @@ according to `SemVer <http://semver.org>`__. The following changes are
 noteworthy:
 
 * The :cpp:class:`nb::enum_\<T\>() <enum_>` binding declaration is now a
-  wrapper that creates either a ``enum.Enum`` or ``enum.IntEnum``-derived type.
+  wrapper that creates either a ``enum.Enum``, ``enum.IntEnum`` or ``enum.Flag``-derived type.
   Previously, nanobind relied on a custom enumeration base class that was a
   frequent source of friction for users.
+  A new flag :cpp:class:`nb::is_flag_enum() <is_flag_enum>`
+  creates a ``enum.Flag``-derived type.
 
   This change may break code that casts entries to integers, which now only
   works for arithmetic (``enum.IntEnum``-derived) enumerations. Replace

docs/classes.rst

@@ -295,8 +295,11 @@ C++11-style strongly typed enumerations.
 
    When the annotation :cpp:class:`nb::is_arithmetic() <is_arithmetic>` is
    passed to :cpp:class:`nb::enum_\<T\> <enum_>`, the resulting Python type
-   will support arithmetic and bit-level operations like comparisons, and, or,
-   xor, negation, etc.
+   will support arithmetic and bit-level operations (and, or,
+   xor, negation).
+   Passing the annotation :cpp:class:`nb::is_flag_enum() <is_flag_enum>` to
+   to :cpp:class:`nb::enum_\<T\> <enum_>`, will result in a Python type
+   `enum.Flags`, that supports bit operations without losing their `Flag` membership.
 
    .. code-block:: cpp
 

include/nanobind/nb_attr.h

@@ -58,6 +58,7 @@ struct is_method {};
 struct is_implicit {};
 struct is_operator {};
 struct is_arithmetic {};
+struct is_flag_enum {};
 struct is_final {};
 struct is_generic {};
 struct kw_only {};

include/nanobind/nb_class.h

@@ -62,9 +62,12 @@ enum class type_flags : uint32_t {
     is_arithmetic            = (1 << 16),
 
     /// Is the number type underlying the enumeration signed?
-    is_signed                = (1 << 17)
+    is_signed                = (1 << 17),
 
-    // One more flag bits available (18) without needing
+    /// Is the underlying enumeration type Flag?
+    is_flag_enum                  = (1 << 18)
+
+    // No more flag bits available (18). Needs
     // a larger reorganization
 };
 
@@ -201,6 +204,10 @@ NB_INLINE void enum_extra_apply(enum_init_data &e, is_arithmetic) {
     e.flags |= (uint32_t) type_flags::is_arithmetic;
 }
 
+NB_INLINE void enum_extra_apply(enum_init_data &e, is_flag_enum) {
+    e.flags |= (uint32_t) type_flags::is_flag_enum;
+}
+
 NB_INLINE void enum_extra_apply(enum_init_data &e, const char *doc) {
     e.docstr = doc;
 }

include/nanobind/nb_python.h

@@ -20,6 +20,7 @@
 
 #include <Python.h>
 #include <frameobject.h>
+#include <object.h>
 #include <pythread.h>
 #include <structmember.h>
 

src/nb_enum.cpp

@@ -1,5 +1,5 @@
 #include "nb_internals.h"
-
+#include <string>
 NAMESPACE_BEGIN(NB_NAMESPACE)
 NAMESPACE_BEGIN(detail)
 
@@ -28,7 +28,7 @@ PyObject *enum_create(enum_init_data *ed) noexcept {
     handle scope(ed->scope);
 
     bool is_arithmetic = ed->flags & (uint32_t) type_flags::is_arithmetic;
-
+    bool is_flag_enum = ed->flags & (uint32_t) type_flags::is_flag_enum;
     str name(ed->name), qualname = name;
     object modname;
 
@@ -43,7 +43,7 @@ PyObject *enum_create(enum_init_data *ed) noexcept {
                 PyUnicode_FromFormat("%U.%U", scope_qualname.ptr(), name.ptr()));
     }
 
-    const char *factory_name = is_arithmetic ? "IntEnum" : "Enum";
+    const char *factory_name = (is_arithmetic || is_flag_enum)  ? (is_flag_enum ? "Flag" : "IntEnum") : "Enum";
 
     object enum_mod = module_::import_("enum"),
            factory = enum_mod.attr(factory_name),
@@ -56,7 +56,11 @@ PyObject *enum_create(enum_init_data *ed) noexcept {
 
     if (is_arithmetic)
         result.attr("__str__") = enum_mod.attr("Enum").attr("__str__");
-
+    if (is_flag_enum)
+        result.attr("__str__") = enum_mod.attr("Flag").attr("__str__");
+#if PY_VERSION_HEX >= 0x030B0000
+    result.attr("_boundary_") =  enum_mod.attr("FlagBoundary").attr("KEEP");
+#endif
     result.attr("__repr__") = result.attr("__str__");
 
     type_init_data *t = new type_init_data();
@@ -150,6 +154,28 @@ bool enum_from_python(const std::type_info *tp, PyObject *o, int64_t *out, uint8
     if (!t)
         return false;
 
+    if ((t->flags & (uint32_t) type_flags::is_flag_enum) !=0) {
+        auto base = PyObject_GetAttrString((PyObject *)o->ob_type, "__base__");
+        auto basename = PyObject_GetAttrString(base, "__name__");
+        Py_ssize_t size;
+        const char* data = PyUnicode_AsUTF8AndSize(basename, &size);
+        std::string s(data, size);
+        if(s == "Flag") {
+            auto pValue = PyObject_GetAttrString(o, "value");
+            if (pValue == nullptr) {
+                PyErr_Clear();
+                return false;
+            }
+            long long value = PyLong_AsLongLong(pValue);
+            if (value == -1 && PyErr_Occurred()) {
+                PyErr_Clear();
+                return false;
+            }
+            *out = (int64_t) value;
+            return true;
+        }
+    }
+
     enum_map *rev = (enum_map *) t->enum_tbl.rev;
     enum_map::iterator it = rev->find((int64_t) (uintptr_t) o);
 
@@ -175,7 +201,6 @@ bool enum_from_python(const std::type_info *tp, PyObject *o, int64_t *out, uint8
         } else {
             unsigned long long value = PyLong_AsUnsignedLongLong(o);
             if (value == (unsigned long long) -1 && PyErr_Occurred()) {
-                PyErr_Clear();
                 return false;
             }
             enum_map::iterator it2 = fwd->find((int64_t) value);
@@ -186,17 +211,19 @@ bool enum_from_python(const std::type_info *tp, PyObject *o, int64_t *out, uint8
         }
 
     }
-
     return false;
 }
 
 PyObject *enum_from_cpp(const std::type_info *tp, int64_t key) noexcept {
     type_data *t = nb_type_c2p(internals, tp);
     if (!t)
         return nullptr;
-
     enum_map *fwd = (enum_map *) t->enum_tbl.fwd;
-
+    if(t->flags & (uint32_t) type_flags::is_flag_enum) {
+        PyObject *value = PyLong_FromLongLong(key);
+        Py_INCREF(value);
+        return value;
+    }
     enum_map::iterator it = fwd->find(key);
     if (it != fwd->end()) {
         PyObject *value = (PyObject *) it->second;

tests/test_enum.cpp

@@ -3,6 +3,7 @@
 namespace nb = nanobind;
 
 enum class Enum  : uint32_t { A, B, C = (uint32_t) -1 };
+enum class Flag  : uint32_t { A = 1, B = 2, C = 4};
 enum class SEnum : int32_t { A, B, C = (int32_t) -1 };
 enum ClassicEnum { Item1, Item2 };
 
@@ -14,6 +15,12 @@ NB_MODULE(test_enum_ext, m) {
         .value("B", Enum::B, "Value B")
         .value("C", Enum::C, "Value C");
 
+    nb::enum_<Flag>(m, "Flag", "enum-level docstring", nb::is_flag_enum())
+            .value("A", Flag::A, "Value A")
+            .value("B", Flag::B, "Value B")
+            .value("C", Flag::C, "Value C")
+            .export_values();
+
     nb::enum_<SEnum>(m, "SEnum", nb::is_arithmetic())
         .value("A", SEnum::A)
         .value("B", SEnum::B)
@@ -31,11 +38,17 @@ NB_MODULE(test_enum_ext, m) {
 
     m.def("from_enum", [](Enum value) { return (uint32_t) value; }, nb::arg().noconvert());
     m.def("to_enum", [](uint32_t value) { return (Enum) value; });
+    m.def("from_enum", [](Flag value) { return (uint32_t) value; }, nb::arg().noconvert());
+    m.def("to_enum", [](uint32_t value) { return (Flag) value; });
     m.def("from_enum", [](SEnum value) { return (int32_t) value; }, nb::arg().noconvert());
-
     m.def("from_enum_implicit", [](Enum value) { return (uint32_t) value; });
 
     m.def("from_enum_default_0", [](Enum value) { return (uint32_t) value; }, nb::arg("value") = Enum::A);
+
+    m.def("from_enum_implicit", [](Flag value) { return (uint32_t) value; });
+
+    m.def("from_enum_default_0", [](Flag value) { return (uint32_t) value; }, nb::arg("value") = Enum::A);
+
     m.def("from_enum_default_1", [](SEnum value) { return (uint32_t) value; }, nb::arg("value") = SEnum::A);
 
     // test for issue #39

tests/test_enum.py

@@ -1,4 +1,5 @@
 import test_enum_ext as t
+import enum
 import pytest
 
 def test01_unsigned_enum():
@@ -136,6 +137,14 @@ def test08_enum_comparisons():
     assert not (t.Enum.B == t.SEnum.B) and not (t.SEnum.B == t.Enum.B)
     assert t.Enum.B != t.SEnum.C and t.SEnum.C != t.Enum.B
 
+def test06_enum_flag():
+    assert (t.Flag(1) | t.Flag(2)).value == 3
+    assert (t.Flag(3) & t.Flag(1)).value == 1
+    assert (t.Flag(3) ^ t.Flag(1)).value == 2
+    assert (t.Flag(3) ==  (t.Flag.A | t.Flag.B))
+    assert (t.from_enum(t.Flag.A | t.Flag.C) == 5)
+    assert (t.from_enum_implicit(t.Flag(1) | t.Flag(4)) == 5)
+
 def test09_enum_methods():
     assert t.Item1.my_value == 0 and t.Item2.my_value == 1
     assert t.Item1.get_value() == 0 and t.Item2.get_value() == 1

tests/test_enum_ext.pyi.ref

@@ -2,6 +2,12 @@ import enum
 from typing import overload
 
 
+A: Flag = Flag.A
+
+B: Flag = Flag.B
+
+C: Flag = Flag.C
+
 class ClassicEnum(enum.Enum):
     Item1 = 0
 
@@ -35,6 +41,18 @@ class EnumProperty:
     @property
     def read_enum(self) -> Enum: ...
 
+class Flag(enum.Flag):
+    """enum-level docstring"""
+
+    A = 1
+    """Value A"""
+
+    B = 2
+    """Value B"""
+
+    C = 4
+    """Value C"""
+
 Item1: ClassicEnum = ClassicEnum.Item1
 
 Item2: ClassicEnum = ClassicEnum.Item2
@@ -49,13 +67,28 @@ class SEnum(enum.IntEnum):
 @overload
 def from_enum(arg: Enum) -> int: ...
 
+@overload
+def from_enum(arg: Flag) -> int: ...
+
 @overload
 def from_enum(arg: SEnum) -> int: ...
 
+@overload
 def from_enum_default_0(value: Enum = Enum.A) -> int: ...
 
+@overload
+def from_enum_default_0(value: Flag = Enum.A) -> int: ...
+
 def from_enum_default_1(value: SEnum = SEnum.A) -> int: ...
 
+@overload
 def from_enum_implicit(arg: Enum, /) -> int: ...
 
+@overload
+def from_enum_implicit(arg: Flag, /) -> int: ...
+
+@overload
 def to_enum(arg: int, /) -> Enum: ...
+
+@overload
+def to_enum(arg: int, /) -> Flag: ...