Add support for pybind11 enum docstrings

sizmailov · Jun 3, 2024 · fe9f906 · fe9f906
1 parent cf58aa6
commit fe9f906
Show file tree

Hide file tree

Showing 11 changed files with 132 additions and 6 deletions.
diff --git a/pybind11_stubgen/parser/mixins/parse.py b/pybind11_stubgen/parser/mixins/parse.py
@@ -171,10 +171,17 @@ def handle_alias(self, path: QualifiedName, origin: Any) -> Alias | None:
         )
 
     def handle_attribute(self, path: QualifiedName, value: Any) -> Attribute | None:
+        doc = None
+        entries = getattr(value, "__entries", None)
+
+        if entries is not None and path[-1] in entries:
+            doc = self.handle_docstring(path, entries[path[-1]][1])
+
         return Attribute(
             name=path[-1],
             value=self.handle_value(value),
             annotation=ResolvedType(name=self.handle_type(type(value))),
+            doc=doc,
         )
 
     def handle_bases(
@@ -195,6 +202,7 @@ def handle_field(self, path: QualifiedName, value: Any) -> Field | None:
             attribute=Attribute(
                 name=attr.name,
                 value=attr.value,
+                doc=attr.doc
             ),
             modifier="static",
         )

diff --git a/pybind11_stubgen/printer.py b/pybind11_stubgen/printer.py
@@ -51,7 +51,10 @@ def print_attribute(self, attr: Attribute) -> list[str]:
             if attr.value is not None:
                 parts.append(f"  # value = {self.print_value(attr.value)}")
 
-        return ["".join(parts)]
+        result = ["".join(parts)]
+        if attr.doc is not None:
+            result.extend(self.print_docstring(attr.doc))
+        return result
 
     def print_argument(self, arg: Argument) -> str:
         parts = []

diff --git a/pybind11_stubgen/structs.py b/pybind11_stubgen/structs.py
@@ -108,6 +108,7 @@ class Attribute:
     name: Identifier
     value: Value | None
     annotation: Annotation | None = field_(default=None)
+    doc: Docstring | None = field_(default=None)
 
 
 @dataclass

diff --git a/stubs/kit-stubs/setup.py b/stubs/kit-stubs/setup.py
@@ -0,0 +1,24 @@
+from setuptools import setup
+import os
+
+
+def find_stubs(package):
+    stubs = []
+    for root, dirs, files in os.walk(package):
+        for file in files:
+            path = os.path.join(root, file).replace(package + os.sep, '', 1)
+            stubs.append(path)
+    return dict(package=stubs)
+
+
+setup(
+    name='kit-stubs',
+    maintainer="kit Developers",
+    maintainer_email="[email protected]",
+    description="PEP 561 type stubs for kit",
+    version='1.0',
+    packages=['kit-stubs'],
+    # PEP 561 requires these
+    install_requires=['kit'],
+    package_data=find_stubs('kit-stubs'),
+)
diff --git a/tests/py-demo/bindings/src/modules/enum.cpp b/tests/py-demo/bindings/src/modules/enum.cpp
@@ -5,11 +5,11 @@
 void bind_enum_module(py::module&&m) {
 
     py::enum_<demo::sublibA::ConsoleForegroundColor>(m, "ConsoleForegroundColor")
-        .value("Green", demo::sublibA::ConsoleForegroundColor::Green)
-        .value("Yellow", demo::sublibA::ConsoleForegroundColor::Yellow)
-        .value("Blue", demo::sublibA::ConsoleForegroundColor::Blue)
-        .value("Magenta", demo::sublibA::ConsoleForegroundColor::Magenta)
-        .value("None_", demo::sublibA::ConsoleForegroundColor::None_)
+        .value("Green", demo::sublibA::ConsoleForegroundColor::Green, "Green color")
+        .value("Yellow", demo::sublibA::ConsoleForegroundColor::Yellow, "Yellow color")
+        .value("Blue", demo::sublibA::ConsoleForegroundColor::Blue, "Blue color")
+        .value("Magenta", demo::sublibA::ConsoleForegroundColor::Magenta, "Magenta color")
+        .value("None_", demo::sublibA::ConsoleForegroundColor::None_, "No color")
         .export_values();
 
     m.def(

diff --git a/tests/stubs/python-3.12/pybind11-master/numpy-array-use-type-var/demo/_bindings/enum.pyi b/tests/stubs/python-3.12/pybind11-master/numpy-array-use-type-var/demo/_bindings/enum.pyi
@@ -30,18 +30,33 @@ class ConsoleForegroundColor:
     Blue: typing.ClassVar[
         ConsoleForegroundColor
     ]  # value = <ConsoleForegroundColor.Blue: 34>
+    """
+    Blue color
+    """
     Green: typing.ClassVar[
         ConsoleForegroundColor
     ]  # value = <ConsoleForegroundColor.Green: 32>
+    """
+    Green color
+    """
     Magenta: typing.ClassVar[
         ConsoleForegroundColor
     ]  # value = <ConsoleForegroundColor.Magenta: 35>
+    """
+    Magenta color
+    """
     None_: typing.ClassVar[
         ConsoleForegroundColor
     ]  # value = <ConsoleForegroundColor.None_: -1>
+    """
+    No color
+    """
     Yellow: typing.ClassVar[
         ConsoleForegroundColor
     ]  # value = <ConsoleForegroundColor.Yellow: 33>
+    """
+    Yellow color
+    """
     __members__: typing.ClassVar[
         dict[str, ConsoleForegroundColor]
     ]  # value = {'Green': <ConsoleForegroundColor.Green: 32>, 'Yellow': <ConsoleForegroundColor.Yellow: 33>, 'Blue': <ConsoleForegroundColor.Blue: 34>, 'Magenta': <ConsoleForegroundColor.Magenta: 35>, 'None_': <ConsoleForegroundColor.None_: -1>}

diff --git a/...stubs/python-3.12/pybind11-master/numpy-array-wrap-with-annotated/demo/_bindings/enum.pyi b/...stubs/python-3.12/pybind11-master/numpy-array-wrap-with-annotated/demo/_bindings/enum.pyi
@@ -30,18 +30,33 @@ class ConsoleForegroundColor:
     Blue: typing.ClassVar[
         ConsoleForegroundColor
     ]  # value = <ConsoleForegroundColor.Blue: 34>
+    """
+    Blue color
+    """
     Green: typing.ClassVar[
         ConsoleForegroundColor
     ]  # value = <ConsoleForegroundColor.Green: 32>
+    """
+    Green color
+    """
     Magenta: typing.ClassVar[
         ConsoleForegroundColor
     ]  # value = <ConsoleForegroundColor.Magenta: 35>
+    """
+    Magenta color
+    """
     None_: typing.ClassVar[
         ConsoleForegroundColor
     ]  # value = <ConsoleForegroundColor.None_: -1>
+    """
+    No color
+    """
     Yellow: typing.ClassVar[
         ConsoleForegroundColor
     ]  # value = <ConsoleForegroundColor.Yellow: 33>
+    """
+    Yellow color
+    """
     __members__: typing.ClassVar[
         dict[str, ConsoleForegroundColor]
     ]  # value = {'Green': <ConsoleForegroundColor.Green: 32>, 'Yellow': <ConsoleForegroundColor.Yellow: 33>, 'Blue': <ConsoleForegroundColor.Blue: 34>, 'Magenta': <ConsoleForegroundColor.Magenta: 35>, 'None_': <ConsoleForegroundColor.None_: -1>}

diff --git a/.../stubs/python-3.12/pybind11-v2.11/numpy-array-wrap-with-annotated/demo/_bindings/enum.pyi b/.../stubs/python-3.12/pybind11-v2.11/numpy-array-wrap-with-annotated/demo/_bindings/enum.pyi
@@ -30,18 +30,33 @@ class ConsoleForegroundColor:
     Blue: typing.ClassVar[
         ConsoleForegroundColor
     ]  # value = <ConsoleForegroundColor.Blue: 34>
+    """
+    Blue color
+    """
     Green: typing.ClassVar[
         ConsoleForegroundColor
     ]  # value = <ConsoleForegroundColor.Green: 32>
+    """
+    Green color
+    """
     Magenta: typing.ClassVar[
         ConsoleForegroundColor
     ]  # value = <ConsoleForegroundColor.Magenta: 35>
+    """
+    Magenta color
+    """
     None_: typing.ClassVar[
         ConsoleForegroundColor
     ]  # value = <ConsoleForegroundColor.None_: -1>
+    """
+    No color
+    """
     Yellow: typing.ClassVar[
         ConsoleForegroundColor
     ]  # value = <ConsoleForegroundColor.Yellow: 33>
+    """
+    Yellow color
+    """
     __members__: typing.ClassVar[
         dict[str, ConsoleForegroundColor]
     ]  # value = {'Green': <ConsoleForegroundColor.Green: 32>, 'Yellow': <ConsoleForegroundColor.Yellow: 33>, 'Blue': <ConsoleForegroundColor.Blue: 34>, 'Magenta': <ConsoleForegroundColor.Magenta: 35>, 'None_': <ConsoleForegroundColor.None_: -1>}

diff --git a/...s/stubs/python-3.12/pybind11-v2.9/numpy-array-wrap-with-annotated/demo/_bindings/enum.pyi b/...s/stubs/python-3.12/pybind11-v2.9/numpy-array-wrap-with-annotated/demo/_bindings/enum.pyi
@@ -30,18 +30,33 @@ class ConsoleForegroundColor:
     Blue: typing.ClassVar[
         ConsoleForegroundColor
     ]  # value = <ConsoleForegroundColor.Blue: 34>
+    """
+    Blue color
+    """
     Green: typing.ClassVar[
         ConsoleForegroundColor
     ]  # value = <ConsoleForegroundColor.Green: 32>
+    """
+    Green color
+    """
     Magenta: typing.ClassVar[
         ConsoleForegroundColor
     ]  # value = <ConsoleForegroundColor.Magenta: 35>
+    """
+    Magenta color
+    """
     None_: typing.ClassVar[
         ConsoleForegroundColor
     ]  # value = <ConsoleForegroundColor.None_: -1>
+    """
+    No color
+    """
     Yellow: typing.ClassVar[
         ConsoleForegroundColor
     ]  # value = <ConsoleForegroundColor.Yellow: 33>
+    """
+    Yellow color
+    """
     __members__: typing.ClassVar[
         dict[str, ConsoleForegroundColor]
     ]  # value = {'Green': <ConsoleForegroundColor.Green: 32>, 'Yellow': <ConsoleForegroundColor.Yellow: 33>, 'Blue': <ConsoleForegroundColor.Blue: 34>, 'Magenta': <ConsoleForegroundColor.Magenta: 35>, 'None_': <ConsoleForegroundColor.None_: -1>}

diff --git a/.../stubs/python-3.7/pybind11-master/numpy-array-wrap-with-annotated/demo/_bindings/enum.pyi b/.../stubs/python-3.7/pybind11-master/numpy-array-wrap-with-annotated/demo/_bindings/enum.pyi
@@ -30,18 +30,33 @@ class ConsoleForegroundColor:
     Blue: typing.ClassVar[
         ConsoleForegroundColor
     ]  # value = <ConsoleForegroundColor.Blue: 34>
+    """
+    Blue color
+    """
     Green: typing.ClassVar[
         ConsoleForegroundColor
     ]  # value = <ConsoleForegroundColor.Green: 32>
+    """
+    Green color
+    """
     Magenta: typing.ClassVar[
         ConsoleForegroundColor
     ]  # value = <ConsoleForegroundColor.Magenta: 35>
+    """
+    Magenta color
+    """
     None_: typing.ClassVar[
         ConsoleForegroundColor
     ]  # value = <ConsoleForegroundColor.None_: -1>
+    """
+    No color
+    """
     Yellow: typing.ClassVar[
         ConsoleForegroundColor
     ]  # value = <ConsoleForegroundColor.Yellow: 33>
+    """
+    Yellow color
+    """
     __members__: typing.ClassVar[
         dict[str, ConsoleForegroundColor]
     ]  # value = {'Green': <ConsoleForegroundColor.Green: 32>, 'Yellow': <ConsoleForegroundColor.Yellow: 33>, 'Blue': <ConsoleForegroundColor.Blue: 34>, 'Magenta': <ConsoleForegroundColor.Magenta: 35>, 'None_': <ConsoleForegroundColor.None_: -1>}

diff --git a/.../stubs/python-3.8/pybind11-master/numpy-array-wrap-with-annotated/demo/_bindings/enum.pyi b/.../stubs/python-3.8/pybind11-master/numpy-array-wrap-with-annotated/demo/_bindings/enum.pyi
@@ -30,18 +30,33 @@ class ConsoleForegroundColor:
     Blue: typing.ClassVar[
         ConsoleForegroundColor
     ]  # value = <ConsoleForegroundColor.Blue: 34>
+    """
+    Blue color
+    """
     Green: typing.ClassVar[
         ConsoleForegroundColor
     ]  # value = <ConsoleForegroundColor.Green: 32>
+    """
+    Green color
+    """
     Magenta: typing.ClassVar[
         ConsoleForegroundColor
     ]  # value = <ConsoleForegroundColor.Magenta: 35>
+    """
+    Magenta color
+    """
     None_: typing.ClassVar[
         ConsoleForegroundColor
     ]  # value = <ConsoleForegroundColor.None_: -1>
+    """
+    No color
+    """
     Yellow: typing.ClassVar[
         ConsoleForegroundColor
     ]  # value = <ConsoleForegroundColor.Yellow: 33>
+    """
+    Yellow color
+    """
     __members__: typing.ClassVar[
         dict[str, ConsoleForegroundColor]
     ]  # value = {'Green': <ConsoleForegroundColor.Green: 32>, 'Yellow': <ConsoleForegroundColor.Yellow: 33>, 'Blue': <ConsoleForegroundColor.Blue: 34>, 'Magenta': <ConsoleForegroundColor.Magenta: 35>, 'None_': <ConsoleForegroundColor.None_: -1>}