Avoid IntoPyObjectExt::into_bound_py_any for None value (#188)

Modifications:
+ Signature-generator can avoid `IntoPyObjectExt::into_bound_py_any` for
`None` value. It's useful for the external types that only impl
FromPyObject, such as cases in relf/egobox#237
  ```rust
  #[gen_stub_pyfunction]
  #[pyfunction(signature = (array=None) )]
  fn my_func(array: Option<numpy::PyReadonlyArray2<f64>>) {}
  ```
+ Add missed newline before function's doc
+ [support Enum's variant
document](cd10de3)
+ [signature support external
repr](a31991a)
When `eval(repr(py_value))` == `py_value`, we can use repr for external
types
  Before:
  ```python
  def default_value(num:Number=...) -> Number: ...
  ```
  After
  ```python
  def default_value(num:Number=Number.FLOAT) -> Number: ...
  ```
Thx!

---------

Co-authored-by: Toshiki Teramura <[email protected]>

Author: zao111222333

examples/pure/pure.pyi

@@ -40,6 +40,9 @@ class Number(Enum):
 
 class NumberRenameAll(Enum):
     FLOAT = ...
+    r"""
+    Float variant
+    """
     INTEGER = ...
 
 def ahash_dict() -> builtins.dict[builtins.str, builtins.int]: ...
@@ -48,17 +51,21 @@ def create_a(x:builtins.int=2) -> A: ...
 
 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.FLOAT) -> Number: ...
 
 def echo_path(path:builtins.str | os.PathLike | pathlib.Path) -> pathlib.Path: ...
 
+def print_c(c:typing.Optional[builtins.int]=None) -> 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.
     """
 

examples/pure/src/lib.rs

@@ -75,6 +75,31 @@ fn create_a(x: usize) -> A {
 #[derive(Debug)]
 struct B;
 
+/// `C` only impl `FromPyObject`
+#[derive(Debug)]
+struct C {
+    x: usize,
+}
+#[gen_stub_pyfunction]
+#[pyfunction(signature = (c=None))]
+fn print_c(c: Option<C>) {
+    if let Some(c) = c {
+        println!("{}", c.x);
+    } else {
+        println!("None");
+    }
+}
+impl FromPyObject<'_> for C {
+    fn extract_bound(ob: &Bound<'_, PyAny>) -> PyResult<Self> {
+        Ok(C { x: ob.extract()? })
+    }
+}
+impl pyo3_stub_gen::PyStubType for C {
+    fn type_output() -> pyo3_stub_gen::TypeInfo {
+        usize::type_output()
+    }
+}
+
 create_exception!(pure, MyError, PyRuntimeError);
 
 /// Returns the length of the string.
@@ -115,6 +140,7 @@ pub enum Number {
 #[pyo3(rename_all = "UPPERCASE")]
 #[derive(Debug, Clone, PartialEq, Eq, Hash)]
 pub enum NumberRenameAll {
+    /// Float variant
     Float,
     Integer,
 }
@@ -159,6 +185,7 @@ fn pure(m: &Bound<PyModule>) -> PyResult<()> {
     m.add_function(wrap_pyfunction!(create_dict, m)?)?;
     m.add_function(wrap_pyfunction!(read_dict, m)?)?;
     m.add_function(wrap_pyfunction!(create_a, m)?)?;
+    m.add_function(wrap_pyfunction!(print_c, m)?)?;
     m.add_function(wrap_pyfunction!(str_len, m)?)?;
     m.add_function(wrap_pyfunction!(echo_path, m)?)?;
     m.add_function(wrap_pyfunction!(ahash_dict, m)?)?;

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

@@ -8,7 +8,7 @@ pub struct PyEnumInfo {
     pyclass_name: String,
     enum_type: Type,
     module: Option<String>,
-    variants: Vec<String>,
+    variants: Vec<(String, String)>,
     doc: String,
 }
 
@@ -54,7 +54,7 @@ impl TryFrom<ItemEnum> for PyEnumInfo {
         let pyclass_name = pyclass_name.unwrap_or_else(|| ident.to_string());
         let variants = variants
             .into_iter()
-            .map(|var| -> Result<String> {
+            .map(|var| -> Result<(String, String)> {
                 let mut var_name = None;
                 for attr in parse_pyo3_attrs(&var.attrs)? {
                     if let Attr::Name(name) = attr {
@@ -65,9 +65,10 @@ impl TryFrom<ItemEnum> for PyEnumInfo {
                 if let Some(renaming_rule) = renaming_rule {
                     var_name = renaming_rule.apply(&var_name);
                 }
-                Ok(var_name)
+                let var_doc = extract_documents(&var.attrs).join("\n");
+                Ok((var_name, var_doc))
             })
-            .collect::<Result<Vec<String>>>()?;
+            .collect::<Result<Vec<(String, String)>>>()?;
         Ok(Self {
             doc,
             enum_type: struct_type,
@@ -88,6 +89,10 @@ impl ToTokens for PyEnumInfo {
             module,
         } = self;
         let module = quote_option(module);
+        let variants: Vec<_> = variants
+            .iter()
+            .map(|(name, doc)| quote! {(#name,#doc)})
+            .collect();
         tokens.append_all(quote! {
             ::pyo3_stub_gen::type_info::PyEnumInfo {
                 pyclass_name: #pyclass_name,

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

@@ -95,22 +95,27 @@ impl ToTokens for ArgsWithSignature<'_> {
                 SignatureArg::Assign(ident, _eq, value) => {
                     let name = ident.to_string();
                     let ty = args_map.get(&name).unwrap();
+                    let default = if value.to_token_stream().to_string() == "None" {
+                        quote! {
+                            "None".to_string()
+                        }
+                    } else {
+                        quote! {
+                            ::pyo3::prepare_freethreaded_python();
+                            ::pyo3::Python::with_gil(|py| -> String {
+                                let v: #ty = #value;
+                                ::pyo3_stub_gen::util::fmt_py_obj(py, v)
+                            })
+                        }
+                    };
                     quote! {
                         ::pyo3_stub_gen::type_info::ArgInfo {
                             name: #name,
                             r#type: <#ty as ::pyo3_stub_gen::PyStubType>::type_input,
                             signature: Some(pyo3_stub_gen::type_info::SignatureArg::Assign{
                                 default: {
                                     static DEFAULT: std::sync::LazyLock<String> = std::sync::LazyLock::new(|| {
-                                        ::pyo3::prepare_freethreaded_python();
-                                        ::pyo3::Python::with_gil(|py| -> String {
-                                            let v: #ty = #value;
-                                            if let Ok(py_obj) = <#ty as ::pyo3::IntoPyObjectExt>::into_bound_py_any(v, py) {
-                                                ::pyo3_stub_gen::util::fmt_py_obj(&py_obj)
-                                            } else {
-                                                "...".to_owned()
-                                            }
-                                        })
+                                        #default
                                     });
                                     &DEFAULT
                                 }

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

@@ -6,7 +6,7 @@ use std::fmt;
 pub struct EnumDef {
     pub name: &'static str,
     pub doc: &'static str,
-    pub variants: &'static [&'static str],
+    pub variants: &'static [(&'static str, &'static str)],
     pub methods: Vec<MethodDef>,
     pub members: Vec<MemberDef>,
 }
@@ -28,8 +28,9 @@ impl fmt::Display for EnumDef {
         writeln!(f, "class {}(Enum):", self.name)?;
         let indent = indent();
         docstring::write_docstring(f, self.doc, indent)?;
-        for variants in self.variants {
-            writeln!(f, "{indent}{} = ...", variants)?;
+        for (variant, variant_doc) in self.variants {
+            writeln!(f, "{indent}{} = ...", variant)?;
+            docstring::write_docstring(f, variant_doc, indent)?;
         }
         for member in &self.members {
             writeln!(f)?;

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

@@ -44,6 +44,7 @@ impl fmt::Display for FunctionDef {
 
         let doc = self.doc;
         if !doc.is_empty() {
+            writeln!(f)?;
             docstring::write_docstring(f, self.doc, indent())?;
         } else {
             writeln!(f, " ...")?;

pyo3-stub-gen/src/type_info.rs

@@ -136,8 +136,8 @@ pub struct PyEnumInfo {
     pub module: Option<&'static str>,
     /// Docstring
     pub doc: &'static str,
-    /// Variants of enum
-    pub variants: &'static [&'static str],
+    /// Variants of enum (name, doc)
+    pub variants: &'static [(&'static str, &'static str)],
 }
 
 inventory::collect!(PyEnumInfo);

pyo3-stub-gen/src/util.rs

@@ -1,4 +1,5 @@
 use pyo3::{prelude::*, types::*};
+use std::ffi::CString;
 
 pub fn all_builtin_types(any: &Bound<'_, PyAny>) -> bool {
     if any.is_instance_of::<PyString>()
@@ -33,19 +34,37 @@ pub fn all_builtin_types(any: &Bound<'_, PyAny>) -> bool {
     false
 }
 
-pub fn fmt_py_obj(any: &Bound<'_, PyAny>) -> String {
-    if all_builtin_types(any) {
-        if let Ok(py_str) = any.repr() {
-            return py_str.to_string();
+/// whether eval(repr(any)) == any
+pub fn valid_external_repr(any: &Bound<'_, PyAny>) -> Option<bool> {
+    let globals = get_globals(any).ok()?;
+    let fmt_str = any.repr().ok()?.to_string();
+    let fmt_cstr = CString::new(fmt_str.clone()).ok()?;
+    let new_any = any.py().eval(&fmt_cstr, Some(&globals), None).ok()?;
+    new_any.eq(any).ok()
+}
+
+fn get_globals<'py>(any: &Bound<'py, PyAny>) -> PyResult<Bound<'py, PyDict>> {
+    let type_object = any.get_type();
+    let type_name = type_object.getattr("__name__")?;
+    let type_name: &str = type_name.extract()?;
+    let globals = PyDict::new(any.py());
+    globals.set_item(type_name, type_object)?;
+    Ok(globals)
+}
+
+pub fn fmt_py_obj<'py, T: pyo3::IntoPyObjectExt<'py>>(py: Python<'py>, obj: T) -> String {
+    if let Ok(any) = obj.into_bound_py_any(py) {
+        if all_builtin_types(&any) || valid_external_repr(&any).is_some_and(|valid| valid) {
+            if let Ok(py_str) = any.repr() {
+                return py_str.to_string();
+            }
         }
     }
     "...".to_owned()
 }
 
 #[cfg(test)]
 mod test {
-    use pyo3::IntoPyObjectExt;
-
     use super::*;
     #[pyclass]
     #[derive(Debug)]
@@ -57,62 +76,69 @@ mod test {
             let dict = PyDict::new(py);
             _ = dict.set_item("k1", "v1");
             _ = dict.set_item("k2", 2);
-            assert_eq!("{'k1': 'v1', 'k2': 2}", fmt_py_obj(&dict));
+            assert_eq!("{'k1': 'v1', 'k2': 2}", fmt_py_obj(py, &dict));
             // class A variable can not be formatted
             _ = dict.set_item("k3", A {});
-            assert_eq!("...", fmt_py_obj(&dict));
+            assert_eq!("...", fmt_py_obj(py, &dict));
         })
     }
     #[test]
     fn test_fmt_list() {
         pyo3::prepare_freethreaded_python();
         Python::with_gil(|py| {
             let list = PyList::new(py, [1, 2]).unwrap();
-            assert_eq!("[1, 2]", fmt_py_obj(&list));
+            assert_eq!("[1, 2]", fmt_py_obj(py, &list));
             // class A variable can not be formatted
             let list = PyList::new(py, [A {}, A {}]).unwrap();
-            assert_eq!("...", fmt_py_obj(&list));
+            assert_eq!("...", fmt_py_obj(py, &list));
         })
     }
     #[test]
     fn test_fmt_tuple() {
         pyo3::prepare_freethreaded_python();
         Python::with_gil(|py| {
             let tuple = PyTuple::new(py, [1, 2]).unwrap();
-            assert_eq!("(1, 2)", fmt_py_obj(&tuple));
+            assert_eq!("(1, 2)", fmt_py_obj(py, tuple));
             let tuple = PyTuple::new(py, [1]).unwrap();
-            assert_eq!("(1,)", fmt_py_obj(&tuple));
+            assert_eq!("(1,)", fmt_py_obj(py, tuple));
             // class A variable can not be formatted
             let tuple = PyTuple::new(py, [A {}]).unwrap();
-            assert_eq!("...", fmt_py_obj(&tuple));
+            assert_eq!("...", fmt_py_obj(py, tuple));
         })
     }
     #[test]
     fn test_fmt_other() {
         pyo3::prepare_freethreaded_python();
         Python::with_gil(|py| {
             // str
-            assert_eq!("'123'", fmt_py_obj(&"123".into_bound_py_any(py).unwrap()));
-            assert_eq!(
-                "\"don't\"",
-                fmt_py_obj(&"don't".into_bound_py_any(py).unwrap())
-            );
-            assert_eq!(
-                "'str\\\\'",
-                fmt_py_obj(&"str\\".into_bound_py_any(py).unwrap())
-            );
+            assert_eq!("'123'", fmt_py_obj(py, &"123"));
+            assert_eq!("\"don't\"", fmt_py_obj(py, &"don't"));
+            assert_eq!("'str\\\\'", fmt_py_obj(py, &"str\\"));
             // bool
-            assert_eq!("True", fmt_py_obj(&true.into_bound_py_any(py).unwrap()));
-            assert_eq!("False", fmt_py_obj(&false.into_bound_py_any(py).unwrap()));
+            assert_eq!("True", fmt_py_obj(py, true));
+            assert_eq!("False", fmt_py_obj(py, false));
             // int
-            assert_eq!("123", fmt_py_obj(&123.into_bound_py_any(py).unwrap()));
+            assert_eq!("123", fmt_py_obj(py, 123));
             // float
-            assert_eq!("1.23", fmt_py_obj(&1.23.into_bound_py_any(py).unwrap()));
+            assert_eq!("1.23", fmt_py_obj(py, 1.23));
             // None
             let none: Option<usize> = None;
-            assert_eq!("None", fmt_py_obj(&none.into_bound_py_any(py).unwrap()));
+            assert_eq!("None", fmt_py_obj(py, none));
             // class A variable can not be formatted
-            assert_eq!("...", fmt_py_obj(&A {}.into_bound_py_any(py).unwrap()));
+            assert_eq!("...", fmt_py_obj(py, A {}));
         })
     }
+    #[test]
+    fn test_fmt_enum() {
+        #[pyclass(eq, eq_int)]
+        #[derive(Debug, Clone, PartialEq, Eq, Hash)]
+        pub enum Number {
+            Float,
+            Integer,
+        }
+        pyo3::prepare_freethreaded_python();
+        Python::with_gil(|py| {
+            assert_eq!("Number.Float", fmt_py_obj(py, Number::Float));
+        });
+    }
 }