Add support for flag enum (#207)

* add support for flag enum

* fix flag selection

* more edge cases

* remove obsolete test and add explanation

src/superqt/combobox/_enum_combobox.py

@@ -1,5 +1,9 @@
-from enum import Enum, EnumMeta
-from typing import Optional, TypeVar
+import sys
+from enum import Enum, EnumMeta, Flag
+from functools import reduce
+from itertools import combinations
+from operator import or_
+from typing import Optional, Tuple, TypeVar
 
 from qtpy.QtCore import Signal
 from qtpy.QtWidgets import QComboBox
@@ -17,10 +21,36 @@ def _get_name(enum_value: Enum):
         # check if function was overloaded
         name = str(enum_value)
     else:
-        name = enum_value.name.replace("_", " ")
+        if enum_value.name is None:
+            # This is hack for python bellow 3.11
+            if not isinstance(enum_value, Flag):
+                raise TypeError(
+                    f"Expected Flag instance, got {enum_value}"
+                )  # pragma: no cover
+            if sys.version_info >= (3, 11):
+                # There is a bug in some releases of Python 3.11 (for example 3.11.3)
+                # that leads to wrong evaluation of or operation on Flag members
+                # and produces numeric value without proper set name property.
+                return f"{enum_value.value}"
+
+            # Before python 3.11 there is no smart name set during
+            # the creation of Flag members.
+            # We needs to decompose the value to get the name.
+            # It is under if condition because it uses private API.
+
+            from enum import _decompose
+
+            members, not_covered = _decompose(enum_value.__class__, enum_value.value)
+            name = "|".join(m.name.replace("_", " ") for m in members[::-1])
+        else:
+            name = enum_value.name.replace("_", " ")
     return name
 
 
+def _get_name_with_value(enum_value: Enum) -> Tuple[str, Enum]:
+    return _get_name(enum_value), enum_value
+
+
 class QEnumComboBox(QComboBox):
     """ComboBox presenting options from a python Enum.
 
@@ -47,9 +77,20 @@ class QEnumComboBox(QComboBox):
         self._allow_none = allow_none and enum is not None
         if allow_none:
             super().addItem(NONE_STRING)
-        names = map(_get_name, self._enum_class.__members__.values())
-        _names = dict.fromkeys(names)  # remove duplicates/aliases, keep order
-        super().addItems(list(_names))
+        names_ = self._get_enum_member_list(enum)
+        super().addItems(list(names_))
+
+    @staticmethod
+    def _get_enum_member_list(enum: Optional[EnumMeta]):
+        if issubclass(enum, Flag):
+            members = list(enum.__members__.values())
+            comb_list = []
+            for i in range(len(members)):
+                comb_list.extend(reduce(or_, x) for x in combinations(members, i + 1))
+
+        else:
+            comb_list = list(enum.__members__.values())
+        return dict(map(_get_name_with_value, comb_list))
 
     def enumClass(self) -> Optional[EnumMeta]:
         """Return current Enum class."""
@@ -70,11 +111,7 @@ class QEnumComboBox(QComboBox):
             if self._allow_none:
                 if self.currentText() == NONE_STRING:
                     return None
-                else:
-                    return list(self._enum_class.__members__.values())[
-                        self.currentIndex() - 1
-                    ]
-            return list(self._enum_class.__members__.values())[self.currentIndex()]
+            return self._get_enum_member_list(self._enum_class)[self.currentText()]
         return None
 
     def setCurrentEnum(self, value: Optional[EnumType]) -> None:

tests/test_enum_comb_box.py

@@ -1,4 +1,5 @@
-from enum import Enum, IntEnum
+import sys
+from enum import Enum, Flag, IntEnum, IntFlag
 
 import pytest
 
@@ -42,6 +43,36 @@ class IntEnum1(IntEnum):
     c = 5
 
 
+class IntFlag1(IntFlag):
+    a = 1
+    b = 2
+    c = 4
+
+
+class Flag1(Flag):
+    a = 1
+    b = 2
+    c = 4
+
+
+class IntFlag2(IntFlag):
+    a = 1
+    b = 2
+    c = 3
+
+
+class Flag2(IntFlag):
+    a = 1
+    b = 2
+    c = 5
+
+
+class FlagOrNum(IntFlag):
+    a = 3
+    b = 5
+    c = 8
+
+
 def test_simple_create(qtbot):
     enum = QEnumComboBox(enum_class=Enum1)
     qtbot.addWidget(enum)
@@ -140,5 +171,60 @@ def test_optional(qtbot):
 def test_simple_create_int_enum(qtbot):
     enum = QEnumComboBox(enum_class=IntEnum1)
     qtbot.addWidget(enum)
-    assert enum.count() == 3
+    assert [enum.itemText(i) for i in range(enum.count())] == ["a", "b", "c"]
+
+
+@pytest.mark.parametrize("enum_class", [IntFlag1, Flag1])
+def test_enum_flag_create(qtbot, enum_class):
+    enum = QEnumComboBox(enum_class=enum_class)
+    qtbot.addWidget(enum)
+    assert [enum.itemText(i) for i in range(enum.count())] == [
+        "a",
+        "b",
+        "c",
+        "a|b",
+        "a|c",
+        "b|c",
+        "a|b|c",
+    ]
+    enum.setCurrentText("a|b")
+    assert enum.currentEnum() == enum_class.a | enum_class.b
+
+
+def test_enum_flag_create_collision(qtbot):
+    enum = QEnumComboBox(enum_class=IntFlag2)
+    qtbot.addWidget(enum)
+    assert [enum.itemText(i) for i in range(enum.count())] == ["a", "b", "c"]
+
+
+@pytest.mark.skipif(
+    sys.version_info >= (3, 11), reason="different representation in 3.11"
+)
+def test_enum_flag_create_collision_evaluated_to_seven(qtbot):
+    enum = QEnumComboBox(enum_class=FlagOrNum)
+    qtbot.addWidget(enum)
+    assert [enum.itemText(i) for i in range(enum.count())] == [
+        "a",
+        "b",
+        "c",
+        "a|b",
+        "a|c",
+        "b|c",
+        "a|b|c",
+    ]
+
+
+@pytest.mark.skipif(
+    sys.version_info < (3, 11), reason="StrEnum is introduced in python 3.11"
+)
+def test_create_str_enum(qtbot):
+    from enum import StrEnum
+
+    class StrEnum1(StrEnum):
+        a = "a"
+        b = "b"
+        c = "c"
+
+    enum = QEnumComboBox(enum_class=StrEnum1)
+    qtbot.addWidget(enum)
     assert [enum.itemText(i) for i in range(enum.count())] == ["a", "b", "c"]