Enums with methods

examples/mixed/python/mixed/main_mod.pyi

@@ -4,18 +4,12 @@
 import builtins
 
 class A:
-    def show_x(self) -> None:
-        ...
-
+    def show_x(self) -> None: ...
 
 class B:
-    def show_x(self) -> None:
-        ...
-
+    def show_x(self) -> None: ...
 
-def create_a(x:builtins.int) -> A:
-    ...
+def create_a(x:builtins.int) -> A: ...
 
-def create_b(x:builtins.int) -> B:
-    ...
+def create_b(x:builtins.int) -> B: ...
 

examples/mixed_sub/python/mixed_sub/main_mod/__init__.pyi

@@ -6,18 +6,12 @@ from . import int
 from . import sub_mod
 
 class A:
-    def show_x(self) -> None:
-        ...
-
+    def show_x(self) -> None: ...
 
 class B:
-    def show_x(self) -> None:
-        ...
-
+    def show_x(self) -> None: ...
 
-def create_a(x:builtins.int) -> A:
-    ...
+def create_a(x:builtins.int) -> A: ...
 
-def create_b(x:builtins.int) -> B:
-    ...
+def create_b(x:builtins.int) -> B: ...
 

examples/mixed_sub/python/mixed_sub/main_mod/int.pyi

@@ -3,6 +3,5 @@
 
 import builtins
 
-def dummy_int_fun(x:builtins.int) -> builtins.int:
-    ...
+def dummy_int_fun(x:builtins.int) -> builtins.int: ...
 

examples/mixed_sub/python/mixed_sub/main_mod/sub_mod.pyi

@@ -4,10 +4,7 @@
 import builtins
 
 class C:
-    def show_x(self) -> None:
-        ...
+    def show_x(self) -> None: ...
 
-
-def create_c(x:builtins.int) -> C:
-    ...
+def create_c(x:builtins.int) -> C: ...
 

examples/mixed_sub_multiple/python/mixed_sub_multiple/main_mod/__init__.pyi

@@ -4,6 +4,5 @@
 from . import mod_a
 from . import mod_b
 
-def greet_main() -> None:
-    ...
+def greet_main() -> None: ...
 

examples/mixed_sub_multiple/python/mixed_sub_multiple/main_mod/mod_a.pyi

@@ -2,6 +2,5 @@
 # ruff: noqa: E501, F401
 
 
-def greet_a() -> None:
-    ...
+def greet_a() -> None: ...
 

examples/mixed_sub_multiple/python/mixed_sub_multiple/main_mod/mod_b.pyi

@@ -2,6 +2,5 @@
 # ruff: noqa: E501, F401
 
 
-def greet_b() -> None:
-    ...
+def greet_b() -> None: ...
 

examples/pure/pure.pyi

@@ -11,46 +11,42 @@ MY_CONSTANT: builtins.int
 class A:
     x: builtins.int
     def __new__(cls,x:builtins.int): ...
-    def show_x(self) -> None:
-        ...
-
-    def ref_test(self, x:dict) -> dict:
-        ...
-
+    def show_x(self) -> None: ...
+    def ref_test(self, x:dict) -> dict: ...
 
 class Number(Enum):
     FLOAT = auto()
     INTEGER = auto()
 
-def ahash_dict() -> builtins.dict[builtins.str, builtins.int]:
-    ...
+    @property
+    def is_float(self) -> builtins.bool:
+        """Whether the number is a float."""
+        ...
+
+    @property
+    def is_integer(self) -> builtins.bool:
+        """Whether the number is an integer."""
+        ...
+
+def ahash_dict() -> builtins.dict[builtins.str, builtins.int]: ...
 
-def create_a(x:builtins.int=2) -> A:
-    ...
+def create_a(x:builtins.int=2) -> A: ...
 
-def create_dict(n:builtins.int) -> builtins.dict[builtins.int, builtins.list[builtins.int]]:
-    ...
+def create_dict(n:builtins.int) -> builtins.dict[builtins.int, builtins.list[builtins.int]]: ...
 
-def default_value(num:Number=...) -> Number:
-    ...
+def default_value(num:Number=...) -> Number: ...
 
-def echo_path(path:builtins.str | os.PathLike | pathlib.Path) -> builtins.str:
-    ...
+def echo_path(path:builtins.str | os.PathLike | pathlib.Path) -> builtins.str: ...
 
-def read_dict(dict:typing.Mapping[builtins.int, typing.Mapping[builtins.int, builtins.int]]) -> None:
-    ...
+def read_dict(dict:typing.Mapping[builtins.int, typing.Mapping[builtins.int, builtins.int]]) -> None: ...
 
-def str_len(x:builtins.str) -> builtins.int:
-    r"""
+def str_len(x:builtins.str) -> builtins.int:    r"""
     Returns the length of the string.
     """
-    ...
 
-def sum(v:typing.Sequence[builtins.int]) -> builtins.int:
-    r"""
+def sum(v:typing.Sequence[builtins.int]) -> builtins.int:    r"""
     Returns the sum of two numbers as a string.
     """
-    ...
 
 class MyError(RuntimeError): ...
 

examples/pure/src/lib.rs

@@ -100,6 +100,22 @@ pub enum Number {
     Integer,
 }
 
+#[gen_stub_pymethods]
+#[pymethods]
+impl Number {
+    #[getter]
+    /// Whether the number is a float.
+    fn is_float(&self) -> bool {
+        matches!(self, Self::Float)
+    }
+
+    #[getter]
+    /// Whether the number is an integer.
+    fn is_integer(&self) -> bool {
+        matches!(self, Self::Integer)
+    }
+}
+
 module_variable!("pure", "MY_CONSTANT", usize);
 
 // Test if non-any PyObject Target can be a default value

pyo3-stub-gen-derive/src/gen_stub.rs

@@ -15,14 +15,17 @@
 //!             MemberInfo {
 //!                 name: "name",
 //!                 r#type: <String as ::pyo3_stub_gen::PyStubType>::type_output,
+//!                 doc: "",
 //!             },
 //!             MemberInfo {
 //!                 name: "ndim",
 //!                 r#type: <usize as ::pyo3_stub_gen::PyStubType>::type_output,
+//!                 doc: "",
 //!             },
 //!             MemberInfo {
 //!                 name: "description",
 //!                 r#type: <Option<String> as ::pyo3_stub_gen::PyStubType>::type_output,
+//!                 doc: "",
 //!             },
 //!         ],
 //!         doc: "",

pyo3-stub-gen-derive/src/gen_stub/member.rs

@@ -1,3 +1,5 @@
+use crate::gen_stub::extract_documents;
+
 use super::{escape_return_type, parse_pyo3_attrs, Attr};
 
 use proc_macro2::TokenStream as TokenStream2;
@@ -6,6 +8,7 @@ use syn::{Error, Field, ImplItemFn, Result, Type};
 
 #[derive(Debug)]
 pub struct MemberInfo {
+    doc: String,
     name: String,
     r#type: Type,
 }
@@ -29,10 +32,12 @@ impl TryFrom<ImplItemFn> for MemberInfo {
     fn try_from(item: ImplItemFn) -> Result<Self> {
         assert!(Self::is_candidate_item(&item)?);
         let ImplItemFn { attrs, sig, .. } = &item;
+        let doc = extract_documents(attrs).join("\n");
         let attrs = parse_pyo3_attrs(attrs)?;
         for attr in attrs {
             if let Attr::Getter(name) = attr {
                 return Ok(MemberInfo {
+                    doc,
                     name: name.unwrap_or(sig.ident.to_string()),
                     r#type: escape_return_type(&sig.output).expect("Getter must return a type"),
                 });
@@ -54,21 +59,28 @@ impl TryFrom<Field> for MemberInfo {
                 field_name = Some(name);
             }
         }
+        let doc = extract_documents(&attrs).join("\n");
         Ok(Self {
             name: field_name.unwrap_or(ident.unwrap().to_string()),
             r#type: ty,
+            doc,
         })
     }
 }
 
 impl ToTokens for MemberInfo {
     fn to_tokens(&self, tokens: &mut TokenStream2) {
-        let Self { name, r#type: ty } = self;
+        let Self {
+            name,
+            r#type: ty,
+            doc,
+        } = self;
         let name = name.strip_prefix("get_").unwrap_or(name);
         tokens.append_all(quote! {
             ::pyo3_stub_gen::type_info::MemberInfo {
                 name: #name,
-                r#type: <#ty as ::pyo3_stub_gen::PyStubType>::type_output
+                r#type: <#ty as ::pyo3_stub_gen::PyStubType>::type_output,
+                doc: #doc,
             }
         })
     }

pyo3-stub-gen-derive/src/gen_stub/pyclass.rs

@@ -124,14 +124,17 @@ mod test {
                 ::pyo3_stub_gen::type_info::MemberInfo {
                     name: "name",
                     r#type: <String as ::pyo3_stub_gen::PyStubType>::type_output,
+                    doc: "",
                 },
                 ::pyo3_stub_gen::type_info::MemberInfo {
                     name: "ndim",
                     r#type: <usize as ::pyo3_stub_gen::PyStubType>::type_output,
+                    doc: "",
                 },
                 ::pyo3_stub_gen::type_info::MemberInfo {
                     name: "description",
                     r#type: <Option<String> as ::pyo3_stub_gen::PyStubType>::type_output,
+                    doc: "",
                 },
             ],
             module: Some("my_module"),

pyo3-stub-gen/src/generate/enum_.rs

@@ -7,6 +7,8 @@ pub struct EnumDef {
     pub name: &'static str,
     pub doc: &'static str,
     pub variants: &'static [&'static str],
+    pub methods: Vec<MethodDef>,
+    pub members: Vec<MemberDef>,
 }
 
 impl From<&PyEnumInfo> for EnumDef {
@@ -15,6 +17,8 @@ impl From<&PyEnumInfo> for EnumDef {
             name: info.pyclass_name,
             doc: info.doc,
             variants: info.variants,
+            methods: Vec::new(),
+            members: Vec::new(),
         }
     }
 }
@@ -34,6 +38,14 @@ impl fmt::Display for EnumDef {
         for variants in self.variants {
             writeln!(f, "{indent}{} = auto()", variants)?;
         }
+        for member in &self.members {
+            writeln!(f)?;
+            member.fmt(f)?;
+        }
+        for methods in &self.methods {
+            writeln!(f)?;
+            methods.fmt(f)?;
+        }
         writeln!(f)?;
         Ok(())
     }

pyo3-stub-gen/src/generate/function.rs

@@ -40,7 +40,7 @@ impl fmt::Display for FunctionDef {
                 write!(f, ", ")?;
             }
         }
-        writeln!(f, ") -> {}:", self.r#return)?;
+        write!(f, ") -> {}:", self.r#return)?;
 
         let doc = self.doc;
         let indent = indent();
@@ -50,8 +50,9 @@ impl fmt::Display for FunctionDef {
                 writeln!(f, "{indent}{}", line)?;
             }
             writeln!(f, r#"{indent}""""#)?;
+        } else {
+            writeln!(f, " ...")?;
         }
-        writeln!(f, "{indent}...")?;
         writeln!(f)?;
         Ok(())
     }

pyo3-stub-gen/src/generate/member.rs

@@ -1,11 +1,18 @@
+use itertools::Itertools;
+
 use crate::{generate::*, type_info::*, TypeInfo};
-use std::{collections::HashSet, fmt};
+use std::{
+    collections::HashSet,
+    fmt::{self},
+};
 
 /// Definition of a class member.
 #[derive(Debug, Clone, PartialEq)]
 pub struct MemberDef {
+    pub is_property: bool,
     pub name: &'static str,
     pub r#type: TypeInfo,
+    pub doc: &'static str,
 }
 
 impl Import for MemberDef {
@@ -17,15 +24,31 @@ impl Import for MemberDef {
 impl From<&MemberInfo> for MemberDef {
     fn from(info: &MemberInfo) -> Self {
         Self {
+            is_property: false,
             name: info.name,
             r#type: (info.r#type)(),
+            doc: info.doc,
         }
     }
 }
 
 impl fmt::Display for MemberDef {
     fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
         let indent = indent();
-        writeln!(f, "{indent}{}: {}", self.name, self.r#type)
+        let doc = self.doc.split("\n").join(&format!("\n{indent}"));
+        if self.is_property {
+            writeln!(f, "{indent}@property")?;
+            writeln!(f, "{indent}def {}(self) -> {}:", self.name, self.r#type)?;
+            if !doc.is_empty() {
+                writeln!(f, r#"{indent}    """{doc}""""#)?;
+            }
+            writeln!(f, "{indent}    ...")
+        } else {
+            writeln!(f, "{indent}{}: {}", self.name, self.r#type)?;
+            if !doc.is_empty() {
+                writeln!(f, r#"{indent}"""{doc}""""#)?;
+            }
+            Ok(())
+        }
     }
 }