Improve default values for str, numbers and bool in text_signature

PyO3 · Mar 16, 2023 · 911e0e3 · 911e0e3
1 parent 658bc6b
commit 911e0e3
Show file tree

Hide file tree

Showing 3 changed files with 44 additions and 8 deletions.
diff --git a/newsfragments/3050.changed.md b/newsfragments/3050.changed.md
@@ -0,0 +1 @@
+Improve default values for str, numbers and bool in automatically-generated `text_signature`.
diff --git a/pyo3-macros-backend/src/pyfunction/signature.rs b/pyo3-macros-backend/src/pyfunction/signature.rs
@@ -585,6 +585,29 @@ impl<'a> FunctionSignature<'a> {
         }
     }
 
+    fn default_value_for_parameter(&self, parameter: &str) -> String {
+        let mut default = "...".to_string();
+        if let Some(fn_arg) = self.arguments.iter().find(|arg| arg.name == parameter) {
+            if let Some(syn::Expr::Lit(syn::ExprLit { lit, .. })) = fn_arg.default.as_ref() {
+                match lit {
+                    syn::Lit::Str(s) => default = s.token().to_string(),
+                    syn::Lit::Char(c) => default = c.token().to_string(),
+                    syn::Lit::Int(i) => default = i.base10_digits().to_string(),
+                    syn::Lit::Float(f) => default = f.base10_digits().to_string(),
+                    syn::Lit::Bool(b) => {
+                        default = if b.value() {
+                            "True".to_string()
+                        } else {
+                            "False".to_string()
+                        }
+                    }
+                    _ => {}
+                }
+            }
+        }
+        default
+    }
+
     pub fn text_signature(&self, fn_type: &FnType) -> String {
         // automatic text signature generation
         let self_argument = match fn_type {
@@ -624,8 +647,8 @@ impl<'a> FunctionSignature<'a> {
             output.push_str(parameter);
 
             if i >= py_sig.required_positional_parameters {
-                // has a default, just use ... for now
-                output.push_str("=...");
+                output.push('=');
+                output.push_str(&self.default_value_for_parameter(parameter));
             }
 
             if py_sig.positional_only_parameters > 0 && i + 1 == py_sig.positional_only_parameters {
@@ -646,8 +669,8 @@ impl<'a> FunctionSignature<'a> {
             maybe_push_comma(&mut output);
             output.push_str(parameter);
             if !required {
-                // has a default, just use ... for now
-                output.push_str("=...")
+                output.push('=');
+                output.push_str(&self.default_value_for_parameter(parameter));
             }
         }
 

diff --git a/tests/test_text_signature.rs b/tests/test_text_signature.rs
@@ -147,6 +147,11 @@ fn test_auto_test_signature_function() {
         let _ = (a, b, args, c, d, kwargs);
     }
 
+    #[pyfunction(signature = (a = 1, /, b = None, c = 1.5, d=5, e = "pyo3", f = 'f', h = true))]
+    fn my_function_5(a: i32, b: Option<i32>, c: f32, d: i32, e: &str, f: char, h: bool) {
+        let _ = (a, b, c, d, e, f, h);
+    }
+
     Python::with_gil(|py| {
         let f = wrap_pyfunction!(my_function)(py).unwrap();
         py_assert!(py, f, "f.__text_signature__ == '(a, b, c)'");
@@ -155,13 +160,20 @@ fn test_auto_test_signature_function() {
         py_assert!(py, f, "f.__text_signature__ == '($module, a, b, c)'");
 
         let f = wrap_pyfunction!(my_function_3)(py).unwrap();
-        py_assert!(py, f, "f.__text_signature__ == '(a, /, b=..., *, c=...)'");
+        py_assert!(py, f, "f.__text_signature__ == '(a, /, b=..., *, c=5)'");
 
         let f = wrap_pyfunction!(my_function_4)(py).unwrap();
         py_assert!(
             py,
             f,
-            "f.__text_signature__ == '(a, /, b=..., *args, c, d=..., **kwargs)'"
+            "f.__text_signature__ == '(a, /, b=..., *args, c, d=5, **kwargs)'"
+        );
+
+        let f = wrap_pyfunction!(my_function_5)(py).unwrap();
+        py_assert!(
+            py,
+            f,
+            "f.__text_signature__ == '(a=1, /, b=..., c=1.5, d=5, e=\"pyo3\", f=\\'f\\', h=True)', f.__text_signature__"
         );
     });
 }
@@ -216,12 +228,12 @@ fn test_auto_test_signature_method() {
         py_assert!(
             py,
             cls,
-            "cls.method_2.__text_signature__ == '($self, a, /, b=..., *, c=...)'"
+            "cls.method_2.__text_signature__ == '($self, a, /, b=..., *, c=5)'"
         );
         py_assert!(
             py,
             cls,
-            "cls.method_3.__text_signature__ == '($self, a, /, b=..., *args, c, d=..., **kwargs)'"
+            "cls.method_3.__text_signature__ == '($self, a, /, b=..., *args, c, d=5, **kwargs)'"
         );
         py_assert!(
             py,