From 96efb0eda9cf4d9774037b779d4a7195685a2b42 Mon Sep 17 00:00:00 2001
From: David Hewitt <1939362+davidhewitt@users.noreply.github.com>
Date: Fri, 3 Feb 2023 06:54:54 +0000
Subject: [PATCH] add better error message for Python in signature

---
 .../src/pyfunction/signature.rs               | 32 +++++++++++++------
 tests/ui/invalid_pyfunction_signatures.rs     |  5 +++
 tests/ui/invalid_pyfunction_signatures.stderr | 10 ++++--
 3 files changed, 35 insertions(+), 12 deletions(-)

diff --git a/pyo3-macros-backend/src/pyfunction/signature.rs b/pyo3-macros-backend/src/pyfunction/signature.rs
index 3a8527458ef..100da063a14 100644
--- a/pyo3-macros-backend/src/pyfunction/signature.rs
+++ b/pyo3-macros-backend/src/pyfunction/signature.rs
@@ -352,10 +352,21 @@ impl<'a> FunctionSignature<'a> {
         let mut parse_state = ParseState::Positional;
         let mut python_signature = PythonSignature::default();
 
-        let mut args_iter = arguments.iter_mut().filter(|arg| !arg.py); // Python<'_> arguments don't show on the Python side.
+        let mut args_iter = arguments.iter_mut();
+
+        let mut next_non_py_argument_checked = |name: &syn::Ident| {
+            for fn_arg in args_iter.by_ref() {
+                if fn_arg.py {
+                    // If the user incorrectly tried to include py: Python in the
+                    // signature, give a useful error as a hint.
+                    ensure_spanned!(
+                        name != fn_arg.name,
+                        name.span() => "arguments of type `Python` must not be part of the signature"
+                    );
+                    // Otherwise try next argument.
+                    continue;
+                }
 
-        let mut next_argument_checked = |name: &syn::Ident| match args_iter.next() {
-            Some(fn_arg) => {
                 ensure_spanned!(
                     name == fn_arg.name,
                     name.span() => format!(
@@ -364,17 +375,17 @@ impl<'a> FunctionSignature<'a> {
                         name.unraw(),
                     )
                 );
-                Ok(fn_arg)
+                return Ok(fn_arg);
             }
-            None => bail_spanned!(
+            bail_spanned!(
                 name.span() => "signature entry does not have a corresponding function argument"
-            ),
+            )
         };
 
         for item in &attribute.value.items {
             match item {
                 SignatureItem::Argument(arg) => {
-                    let fn_arg = next_argument_checked(&arg.ident)?;
+                    let fn_arg = next_non_py_argument_checked(&arg.ident)?;
                     parse_state.add_argument(
                         &mut python_signature,
                         arg.ident.unraw().to_string(),
@@ -389,12 +400,12 @@ impl<'a> FunctionSignature<'a> {
                     parse_state.finish_pos_args(&python_signature, sep.span())?
                 }
                 SignatureItem::Varargs(varargs) => {
-                    let fn_arg = next_argument_checked(&varargs.ident)?;
+                    let fn_arg = next_non_py_argument_checked(&varargs.ident)?;
                     fn_arg.is_varargs = true;
                     parse_state.add_varargs(&mut python_signature, &varargs)?;
                 }
                 SignatureItem::Kwargs(kwargs) => {
-                    let fn_arg = next_argument_checked(&kwargs.ident)?;
+                    let fn_arg = next_non_py_argument_checked(&kwargs.ident)?;
                     fn_arg.is_kwargs = true;
                     parse_state.add_kwargs(&mut python_signature, &kwargs)?;
                 }
@@ -404,7 +415,8 @@ impl<'a> FunctionSignature<'a> {
             };
         }
 
-        if let Some(arg) = args_iter.next() {
+        // Ensure no non-py arguments remain
+        if let Some(arg) = args_iter.find(|arg| !arg.py) {
             bail_spanned!(
                 attribute.kw.span() => format!("missing signature entry for argument `{}`", arg.name)
             );
diff --git a/tests/ui/invalid_pyfunction_signatures.rs b/tests/ui/invalid_pyfunction_signatures.rs
index 2bf849c0502..f5a9bee4e6c 100644
--- a/tests/ui/invalid_pyfunction_signatures.rs
+++ b/tests/ui/invalid_pyfunction_signatures.rs
@@ -44,6 +44,11 @@ fn function_with_kwargs_after_kwargs(kwargs_a: Option<&PyDict>, kwargs_b: Option
     let _ = kwargs_b;
 }
 
+#[pyfunction(signature = (py))]
+fn signature_contains_py(py: Python<'_>) {
+    let _ = py;
+}
+
 #[pyclass]
 struct MyClass;
 
diff --git a/tests/ui/invalid_pyfunction_signatures.stderr b/tests/ui/invalid_pyfunction_signatures.stderr
index 6720bb8753d..44031cceec7 100644
--- a/tests/ui/invalid_pyfunction_signatures.stderr
+++ b/tests/ui/invalid_pyfunction_signatures.stderr
@@ -46,8 +46,14 @@ error: `**kwargs_b` not allowed after `**kwargs_a`
 41 | #[pyo3(signature = (**kwargs_a, **kwargs_b))]
    |                                 ^
 
+error: arguments of type `Python` must not be part of the signature
+  --> tests/ui/invalid_pyfunction_signatures.rs:47:27
+   |
+47 | #[pyfunction(signature = (py))]
+   |                           ^^
+
 error: cannot define both function signature and legacy arguments
-  --> tests/ui/invalid_pyfunction_signatures.rs:53:12
+  --> tests/ui/invalid_pyfunction_signatures.rs:58:12
    |
-53 |     #[pyo3(signature = (x))]
+58 |     #[pyo3(signature = (x))]
    |            ^^^^^^^^^