add native enum support

Author: InvincibleRMC

mypy/stubgenc.py

@@ -870,8 +870,12 @@ 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} = ...")
+            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,23 @@
+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 = ...
+    RED = ...
+    _generate_next_value_: ClassVar[Callable] = ...
+    _member_map_: ClassVar[dict] = ...
+    _member_names_: ClassVar[list] = ...
+    _member_type_: ClassVar[type[object]] = ...
+    _unhashable_values_: ClassVar[list] = ...
+    _use_args_: ClassVar[bool] = ...
+    _value2member_map_: ClassVar[dict] = ...
+    _value_repr_: ClassVar[None] = ...
+
 class Point:
     class AngleUnit:
         __members__: ClassVar[dict] = ...  # read-only

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

@@ -1,9 +1,23 @@
+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 = ...
+    RED = ...
+    _generate_next_value_: ClassVar[Callable] = ...
+    _member_map_: ClassVar[dict] = ...
+    _member_names_: ClassVar[list] = ...
+    _member_type_: ClassVar[type[object]] = ...
+    _unhashable_values_: ClassVar[list] = ...
+    _use_args_: ClassVar[bool] = ...
+    _value2member_map_: ClassVar[dict] = ...
+    _value_repr_: ClassVar[None] = ...
+
 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")
+    .value("RED", Color::RED)
+    .value("GREEN", Color::GREEN)
+    .finalize();
+
 
   pyPoint
     .def(py::init<>())