[stubgenc] Add py::native enum<> support

mypy/stubgenc.py

@@ -14,7 +14,7 @@
 import os.path
 from collections.abc import Mapping
 from types import FunctionType, ModuleType
-from typing import Any, Callable
+from typing import Any, Callable, Final
 
 from mypy.fastparse import parse_type_comment
 from mypy.moduleinspect import is_c_module
@@ -40,6 +40,17 @@
 )
 from mypy.util import quote_docstring
 
+SPECIAL_ENUM_ATTRS: Final = {
+    "_generate_next_value_",
+    "_member_map_",
+    "_member_names_",
+    "_member_type_",
+    "_unhashable_values_",
+    "_use_args_",
+    "_value2member_map_",
+    "_value_repr_",
+}
+
 
 class ExternalSignatureGenerator(SignatureGenerator):
     def __init__(
@@ -862,6 +873,8 @@ def generate_class_stub(
                 )
             elif inspect.isclass(value) and self.is_defined_in_module(value):
                 self.generate_class_stub(attr, value, types, parent_class=class_info)
+            elif attr in SPECIAL_ENUM_ATTRS:
+                pass
             else:
                 attrs.append((attr, value))
 
@@ -870,8 +883,14 @@ def generate_class_stub(
                 # special case for __hash__
                 continue
             prop_type_name = self.strip_or_import(self.get_type_annotation(value))
-            classvar = self.add_name("typing.ClassVar")
-            static_properties.append(f"{self._indent}{attr}: {classvar}[{prop_type_name}] = ...")
+            # Pybind supports native enums now
+            if issubclass(cls, enum.Enum) and attr in cls._member_names_:
+                static_properties.append(f"{self._indent}{attr} = {cls._member_map_[attr].value}")
+            else:
+                classvar = self.add_name("typing.ClassVar")
+                static_properties.append(
+                    f"{self._indent}{attr}: {classvar}[{prop_type_name}] = ..."
+                )
 
         self.dedent()
 

test-data/pybind11_fixtures/expected_stubs_no_docs/pybind11_fixtures/demo.pyi

@@ -1,9 +1,15 @@
+import enum
 import typing
-from typing import ClassVar, overload
+from typing import Callable, ClassVar, overload
 
 PI: float
 __version__: str
 
+class Color(enum.Enum):
+    __new__: ClassVar[Callable] = ...
+    GREEN = 1
+    RED = 0
+
 class Point:
     class AngleUnit:
         __members__: ClassVar[dict] = ...  # read-only

test-data/pybind11_fixtures/expected_stubs_with_docs/pybind11_fixtures/demo.pyi

@@ -1,9 +1,16 @@
+import enum
 import typing
-from typing import ClassVar, overload
+from typing import Callable, ClassVar, overload
 
 PI: float
 __version__: str
 
+class Color(enum.Enum):
+    """Color Enum"""
+    __new__: ClassVar[Callable] = ...
+    GREEN = 1
+    RED = 0
+
 class Point:
     class AngleUnit:
         """Members:

test-data/pybind11_fixtures/src/main.cpp

@@ -51,6 +51,7 @@
 #include <pybind11/pybind11.h>
 #include <pybind11/stl.h>
 #include <pybind11/stl/filesystem.h>
+#include <pybind11/native_enum.h>
 
 namespace py = pybind11;
 
@@ -212,6 +213,11 @@ const Point Point::y_axis = Point(0, 1);
 Point::LengthUnit Point::length_unit = Point::LengthUnit::mm;
 Point::AngleUnit Point::angle_unit = Point::AngleUnit::radian;
 
+enum Color {
+  RED = 0,
+  GREEN = 1
+};
+
 } // namespace: demo
 
 // Bindings
@@ -230,6 +236,11 @@ void bind_demo(py::module& m) {
   py::class_<Point> pyPoint(m, "Point");
   py::enum_<Point::LengthUnit> pyLengthUnit(pyPoint, "LengthUnit");
   py::enum_<Point::AngleUnit> pyAngleUnit(pyPoint, "AngleUnit");
+  py::native_enum<Color>(m, "Color", "enum.Enum", "Color Enum")
+    .value("RED", Color::RED)
+    .value("GREEN", Color::GREEN)
+    .finalize();
+
 
   pyPoint
     .def(py::init<>())