diff --git a/mypy/stubgenc.py b/mypy/stubgenc.py index e0e063927aad..6b4c51d3e3a4 100755 --- a/mypy/stubgenc.py +++ b/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() diff --git a/test-data/pybind11_fixtures/expected_stubs_no_docs/pybind11_fixtures/demo.pyi b/test-data/pybind11_fixtures/expected_stubs_no_docs/pybind11_fixtures/demo.pyi index 09e75e1ad4aa..f491683987b8 100644 --- a/test-data/pybind11_fixtures/expected_stubs_no_docs/pybind11_fixtures/demo.pyi +++ b/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 diff --git a/test-data/pybind11_fixtures/expected_stubs_with_docs/pybind11_fixtures/demo.pyi b/test-data/pybind11_fixtures/expected_stubs_with_docs/pybind11_fixtures/demo.pyi index 580aa2700178..400826965087 100644 --- a/test-data/pybind11_fixtures/expected_stubs_with_docs/pybind11_fixtures/demo.pyi +++ b/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: diff --git a/test-data/pybind11_fixtures/src/main.cpp b/test-data/pybind11_fixtures/src/main.cpp index 4d275ab1fd70..2d5291d3c635 100644 --- a/test-data/pybind11_fixtures/src/main.cpp +++ b/test-data/pybind11_fixtures/src/main.cpp @@ -51,6 +51,7 @@ #include #include #include +#include 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_ pyPoint(m, "Point"); py::enum_ pyLengthUnit(pyPoint, "LengthUnit"); py::enum_ pyAngleUnit(pyPoint, "AngleUnit"); + py::native_enum(m, "Color", "enum.Enum", "Color Enum") + .value("RED", Color::RED) + .value("GREEN", Color::GREEN) + .finalize(); + pyPoint .def(py::init<>())