Adds basic input type annotations

Author: Tpt

guide/src/class.md

@@ -1378,6 +1378,9 @@ impl pyo3::types::DerefToPyAny for MyClass {}
 unsafe impl pyo3::type_object::PyTypeInfo for MyClass {
     const NAME: &'static str = "MyClass";
     const MODULE: ::std::option::Option<&'static str> = ::std::option::Option::None;
+    #[cfg(feature = "experimental-inspect")]
+    const INPUT_TYPE: &'static str = "MyClass";
+
     #[inline]
     fn type_object_raw(py: pyo3::Python<'_>) -> *mut pyo3::ffi::PyTypeObject {
         <Self as pyo3::impl_::pyclass::PyClassImpl>::lazy_type_object()
@@ -1393,6 +1396,8 @@ impl pyo3::PyClass for MyClass {
 impl<'a, 'py> pyo3::impl_::extract_argument::PyFunctionArgument<'a, 'py, false> for &'a MyClass
 {
     type Holder = ::std::option::Option<pyo3::PyRef<'py, MyClass>>;
+    #[cfg(feature = "experimental-inspect")]
+    const INPUT_TYPE: &'static str = "MyClass";
 
     #[inline]
     fn extract(obj: &'a pyo3::Bound<'py, PyAny>, holder: &'a mut Self::Holder) -> pyo3::PyResult<Self> {

noxfile.py

@@ -830,6 +830,7 @@ def update_ui_tests(session: nox.Session):
 @nox.session(name="test-introspection")
 def test_introspection(session: nox.Session):
     session.install("maturin")
+    session.install("ruff")
     target = os.environ.get("CARGO_BUILD_TARGET")
     for options in ([], ["--release"]):
         if target is not None:

pyo3-introspection/Cargo.toml

@@ -14,5 +14,8 @@ goblin = "0.9.0"
 serde = { version = "1", features = ["derive"] }
 serde_json = "1"
 
+[dev-dependencies]
+tempfile = "3.12.0"
+
 [lints]
 workspace = true

pyo3-introspection/src/introspection.rs

@@ -112,6 +112,7 @@ fn convert_argument(arg: &ChunkArgument) -> Argument {
     Argument {
         name: arg.name.clone(),
         default_value: arg.default.clone(),
+        annotation: arg.annotation.clone(),
     }
 }
 
@@ -315,4 +316,6 @@ struct ChunkArgument {
     name: String,
     #[serde(default)]
     default: Option<String>,
+    #[serde(default)]
+    annotation: Option<String>,
 }

pyo3-introspection/src/model.rs

@@ -36,6 +36,8 @@ pub struct Argument {
     pub name: String,
     /// Default value as a Python expression
     pub default_value: Option<String>,
+    /// Type annotation as a Python expression
+    pub annotation: Option<String>,
 }
 
 /// A variable length argument ie. *vararg or **kwarg

pyo3-introspection/src/stubs.rs

@@ -1,5 +1,5 @@
 use crate::model::{Argument, Class, Function, Module, VariableLengthArgument};
-use std::collections::HashMap;
+use std::collections::{BTreeSet, HashMap};
 use std::path::{Path, PathBuf};
 
 /// Generates the [type stubs](https://typing.readthedocs.io/en/latest/source/stubs.html) of a given module.
@@ -32,51 +32,70 @@ fn add_module_stub_files(
 
 /// Generates the module stubs to a String, not including submodules
 fn module_stubs(module: &Module) -> String {
+    let mut modules_to_import = BTreeSet::new();
     let mut elements = Vec::new();
     for class in &module.classes {
         elements.push(class_stubs(class));
     }
     for function in &module.functions {
-        elements.push(function_stubs(function));
+        elements.push(function_stubs(function, &mut modules_to_import));
     }
     elements.push(String::new()); // last line jump
-    elements.join("\n")
+
+    let mut final_elements = Vec::new();
+    for module_to_import in &modules_to_import {
+        final_elements.push(format!("import {module_to_import}"));
+    }
+    final_elements.extend(elements);
+    final_elements.join("\n")
 }
 
 fn class_stubs(class: &Class) -> String {
     format!("class {}: ...", class.name)
 }
 
-fn function_stubs(function: &Function) -> String {
+fn function_stubs(function: &Function, modules_to_import: &mut BTreeSet<String>) -> String {
     // Signature
     let mut parameters = Vec::new();
     for argument in &function.arguments.positional_only_arguments {
-        parameters.push(argument_stub(argument));
+        parameters.push(argument_stub(argument, modules_to_import));
     }
     if !function.arguments.positional_only_arguments.is_empty() {
         parameters.push("/".into());
     }
     for argument in &function.arguments.arguments {
-        parameters.push(argument_stub(argument));
+        parameters.push(argument_stub(argument, modules_to_import));
     }
     if let Some(argument) = &function.arguments.vararg {
         parameters.push(format!("*{}", variable_length_argument_stub(argument)));
     } else if !function.arguments.keyword_only_arguments.is_empty() {
         parameters.push("*".into());
     }
     for argument in &function.arguments.keyword_only_arguments {
-        parameters.push(argument_stub(argument));
+        parameters.push(argument_stub(argument, modules_to_import));
     }
     if let Some(argument) = &function.arguments.kwarg {
         parameters.push(format!("**{}", variable_length_argument_stub(argument)));
     }
     format!("def {}({}): ...", function.name, parameters.join(", "))
 }
 
-fn argument_stub(argument: &Argument) -> String {
+fn argument_stub(argument: &Argument, modules_to_import: &mut BTreeSet<String>) -> String {
     let mut output = argument.name.clone();
+    if let Some(annotation) = &argument.annotation {
+        output.push_str(": ");
+        output.push_str(annotation);
+        if let Some((module, _)) = annotation.rsplit_once('.') {
+            // TODO: this is very naive
+            modules_to_import.insert(module.into());
+        }
+    }
     if let Some(default_value) = &argument.default_value {
-        output.push('=');
+        output.push_str(if argument.annotation.is_some() {
+            " = "
+        } else {
+            "="
+        });
         output.push_str(default_value);
     }
     output
@@ -99,26 +118,29 @@ mod tests {
                 positional_only_arguments: vec![Argument {
                     name: "posonly".into(),
                     default_value: None,
+                    annotation: None,
                 }],
                 arguments: vec![Argument {
                     name: "arg".into(),
                     default_value: None,
+                    annotation: None,
                 }],
                 vararg: Some(VariableLengthArgument {
                     name: "varargs".into(),
                 }),
                 keyword_only_arguments: vec![Argument {
                     name: "karg".into(),
                     default_value: None,
+                    annotation: Some("str".into()),
                 }],
                 kwarg: Some(VariableLengthArgument {
                     name: "kwarg".into(),
                 }),
             },
         };
         assert_eq!(
-            "def func(posonly, /, arg, *varargs, karg, **kwarg): ...",
-            function_stubs(&function)
+            "def func(posonly, /, arg, *varargs, karg: str, **kwarg): ...",
+            function_stubs(&function, &mut BTreeSet::new())
         )
     }
 
@@ -130,22 +152,25 @@ mod tests {
                 positional_only_arguments: vec![Argument {
                     name: "posonly".into(),
                     default_value: Some("1".into()),
+                    annotation: None,
                 }],
                 arguments: vec![Argument {
                     name: "arg".into(),
                     default_value: Some("True".into()),
+                    annotation: None,
                 }],
                 vararg: None,
                 keyword_only_arguments: vec![Argument {
                     name: "karg".into(),
                     default_value: Some("\"foo\"".into()),
+                    annotation: Some("str".into()),
                 }],
                 kwarg: None,
             },
         };
         assert_eq!(
-            "def afunc(posonly=1, /, arg=True, *, karg=\"foo\"): ...",
-            function_stubs(&function)
+            "def afunc(posonly=1, /, arg=True, *, karg: str = \"foo\"): ...",
+            function_stubs(&function, &mut BTreeSet::new())
         )
     }
 }

pyo3-introspection/tests/test.rs

@@ -1,8 +1,11 @@
-use anyhow::Result;
+use anyhow::{ensure, Result};
 use pyo3_introspection::{introspect_cdylib, module_stub_files};
 use std::collections::HashMap;
+use std::io::{Read, Seek, SeekFrom, Write};
 use std::path::{Path, PathBuf};
+use std::process::Command;
 use std::{env, fs};
+use tempfile::NamedTempFile;
 
 #[test]
 fn pytests_stubs() -> Result<()> {
@@ -42,9 +45,12 @@ fn pytests_stubs() -> Result<()> {
                 file_name.display()
             )
         });
+
+        let actual_file_content = format_with_ruff(actual_file_content)?;
+
         assert_eq!(
-            &expected_file_content.replace('\r', ""), // Windows compatibility
-            actual_file_content,
+            expected_file_content.as_str(),
+            actual_file_content.as_str(),
             "The content of file {} is different",
             file_name.display()
         )
@@ -75,3 +81,25 @@ fn add_dir_files(
     }
     Ok(())
 }
+
+fn format_with_ruff(code: &str) -> Result<String> {
+    let temp_file = NamedTempFile::with_suffix(".pyi")?;
+    // Write to file
+    {
+        let mut file = temp_file.as_file();
+        file.write_all(code.as_bytes())?;
+        file.flush()?;
+        file.seek(SeekFrom::Start(0))?;
+    }
+    ensure!(
+        Command::new("ruff")
+            .arg("format")
+            .arg(temp_file.path())
+            .status()?
+            .success(),
+        "Failed to run ruff"
+    );
+    let mut content = String::new();
+    temp_file.as_file().read_to_string(&mut content)?;
+    Ok(content)
+}