New TypeInfo struct and PyStubType trait (#37)

- Resolve #16, remove the dependency to "PyO3/experimental-inspect" - New `TypeInfo` is introduced which replace `pyo3::inspect::types::TypeInfo`. It is generated by `PyStubType` traits. - Impelement `PyStubType` to - builtins (`u32`, `i32`, `String`, ...) - Collections (`Vec<T>`, `HashMap<K, V>`, ...) - PyO3 primitives (`PyAny`, `PyList`, `PyDict`) - Helper for implementing `PyStubType` for user struct, i.e. `#[derive(PyStubType)]` will be next PR.
Jij-Inc · Aug 14, 2024 · adc4c54 · adc4c54
1 parent 7e41b75
commit adc4c54
Show file tree

Hide file tree

Showing 29 changed files with 471 additions and 102 deletions.
diff --git a/Cargo.lock b/Cargo.lock
diff --git a/Cargo.toml b/Cargo.toml
@@ -8,7 +8,7 @@ members = [
 resolver = "2"
 
 [workspace.package]
-version = "0.4.1"
+version = "0.5.0"
 edition = "2021"
 
 description = "Stub file (*.pyi) generator for PyO3"
@@ -24,9 +24,10 @@ insta = "1.39.0"
 inventory = "0.3.15"
 itertools = "0.13.0"
 log = "0.4.22"
+maplit = "1.0.2"
 prettyplease = "0.2.20"
 proc-macro2 = "1.0.86"
-pyo3 = { version = ">=0.21.0", features = ["experimental-inspect"] }
+pyo3 = ">= 0.21.0"
 quote = "1.0.36"
 serde = { version = "1.0.206", features = ["derive"] }
 syn = "2.0.74"

diff --git a/pyo3-stub-gen-derive/src/gen_stub.rs b/pyo3-stub-gen-derive/src/gen_stub.rs
@@ -14,15 +14,15 @@
 //!         members: &[
 //!             MemberInfo {
 //!                 name: "name",
-//!                 r#type: <String as IntoPy<PyObject>>::type_output,
+//!                 r#type: <String as ::pyo3_stub_gen::PyStubType>::type_output,
 //!             },
 //!             MemberInfo {
 //!                 name: "ndim",
-//!                 r#type: <usize as IntoPy<PyObject>>::type_output,
+//!                 r#type: <usize as ::pyo3_stub_gen::PyStubType>::type_output,
 //!             },
 //!             MemberInfo {
 //!                 name: "description",
-//!                 r#type: <Option<String> as IntoPy<PyObject>>::type_output,
+//!                 r#type: <Option<String> as ::pyo3_stub_gen::PyStubType>::type_output,
 //!             },
 //!         ],
 //!         doc: "",

diff --git a/pyo3-stub-gen-derive/src/gen_stub/arg.rs b/pyo3-stub-gen-derive/src/gen_stub/arg.rs
@@ -67,7 +67,7 @@ fn type_to_token(ty: &Type) -> TokenStream2 {
                 if last_seg.ident == "CompareOp" {
                     quote! { ::pyo3_stub_gen::type_info::compare_op_type_input }
                 } else {
-                    quote! { <#ty as FromPyObject>::type_input }
+                    quote! { <#ty as ::pyo3_stub_gen::PyStubType>::type_input }
                 }
             } else {
                 unreachable!("Empty path segment: {:?}", path);
@@ -78,24 +78,24 @@ fn type_to_token(ty: &Type) -> TokenStream2 {
                 Type::Path(TypePath { path, .. }) => {
                     if let Some(last) = path.segments.last() {
                         match last.ident.to_string().as_str() {
-                            // Types where `&T: FromPyObject` instead of `T: FromPyObject`
+                            // Types where `&T: ::pyo3_stub_gen::PyStubType` instead of `T: ::pyo3_stub_gen::PyStubType`
                             // i.e. `&str` and most of `Py*` types defined in PyO3.
                             "str" | "PyAny" | "PyString" | "PyDict" => {
-                                return quote! { <#ty as FromPyObject>::type_input };
+                                return quote! { <#ty as ::pyo3_stub_gen::PyStubType>::type_input };
                             }
                             _ => {}
                         }
                     }
                 }
                 Type::Slice(_) => {
-                    return quote! { <#ty as FromPyObject>::type_input };
+                    return quote! { <#ty as ::pyo3_stub_gen::PyStubType>::type_input };
                 }
                 _ => {}
             }
             type_to_token(elem)
         }
         _ => {
-            quote! { <#ty as FromPyObject>::type_input }
+            quote! { <#ty as ::pyo3_stub_gen::PyStubType>::type_input }
         }
     }
 }

diff --git a/pyo3-stub-gen-derive/src/gen_stub/member.rs b/pyo3-stub-gen-derive/src/gen_stub/member.rs
@@ -68,7 +68,7 @@ impl ToTokens for MemberInfo {
         tokens.append_all(quote! {
             ::pyo3_stub_gen::type_info::MemberInfo {
                 name: #name,
-                r#type: <#ty as IntoPy<PyObject>>::type_output
+                r#type: <#ty as ::pyo3_stub_gen::PyStubType>::type_output
             }
         })
     }

diff --git a/pyo3-stub-gen-derive/src/gen_stub/method.rs b/pyo3-stub-gen-derive/src/gen_stub/method.rs
@@ -100,7 +100,7 @@ impl ToTokens for MethodInfo {
         } = self;
         let sig_tt = quote_option(sig);
         let ret_tt = if let Some(ret) = ret {
-            quote! { <#ret as IntoPy<::pyo3::PyObject>>::type_output }
+            quote! { <#ret as pyo3_stub_gen::PyStubType>::type_output }
         } else {
             quote! { ::pyo3_stub_gen::type_info::no_return_type_output }
         };

diff --git a/pyo3-stub-gen-derive/src/gen_stub/pyclass.rs b/pyo3-stub-gen-derive/src/gen_stub/pyclass.rs
@@ -107,15 +107,15 @@ mod test {
             members: &[
                 ::pyo3_stub_gen::type_info::MemberInfo {
                     name: "name",
-                    r#type: <String as IntoPy<PyObject>>::type_output,
+                    r#type: <String as ::pyo3_stub_gen::PyStubType>::type_output,
                 },
                 ::pyo3_stub_gen::type_info::MemberInfo {
                     name: "ndim",
-                    r#type: <usize as IntoPy<PyObject>>::type_output,
+                    r#type: <usize as ::pyo3_stub_gen::PyStubType>::type_output,
                 },
                 ::pyo3_stub_gen::type_info::MemberInfo {
                     name: "description",
-                    r#type: <Option<String> as IntoPy<PyObject>>::type_output,
+                    r#type: <Option<String> as ::pyo3_stub_gen::PyStubType>::type_output,
                 },
             ],
             module: Some("my_module"),

diff --git a/pyo3-stub-gen-derive/src/gen_stub/pyfunction.rs b/pyo3-stub-gen-derive/src/gen_stub/pyfunction.rs
@@ -84,7 +84,7 @@ impl ToTokens for PyFunctionInfo {
             module,
         } = self;
         let ret_tt = if let Some(ret) = ret {
-            quote! { <#ret as IntoPy<::pyo3::PyObject>>::type_output }
+            quote! { <#ret as pyo3_stub_gen::PyStubType>::type_output }
         } else {
             quote! { ::pyo3_stub_gen::type_info::no_return_type_output }
         };

diff --git a/pyo3-stub-gen-derive/src/lib.rs b/pyo3-stub-gen-derive/src/lib.rs
@@ -5,10 +5,8 @@ use proc_macro::TokenStream;
 /// Embed metadata for Python stub file generation for `#[pyclass]` macro
 ///
 /// ```
-/// # use pyo3_stub_gen_derive::*;
-/// # use pyo3::*;
-/// #[gen_stub_pyclass]
-/// #[pyclass(mapping, module = "my_module", name = "Placeholder")]
+/// #[pyo3_stub_gen_derive::gen_stub_pyclass]
+/// #[pyo3::pyclass(mapping, module = "my_module", name = "Placeholder")]
 /// #[derive(Debug, Clone)]
 /// pub struct PyPlaceholder {
 ///     #[pyo3(get)]
@@ -30,10 +28,8 @@ pub fn gen_stub_pyclass(_attr: TokenStream, item: TokenStream) -> TokenStream {
 /// Embed metadata for Python stub file generation for `#[pyclass]` macro with enum
 ///
 /// ```
-/// # use pyo3_stub_gen_derive::*;
-/// # use pyo3::*;
-/// #[gen_stub_pyclass_enum]
-/// #[pyclass(module = "my_module", name = "DataType")]
+/// #[pyo3_stub_gen_derive::gen_stub_pyclass_enum]
+/// #[pyo3::pyclass(module = "my_module", name = "DataType")]
 /// #[derive(Debug, Clone, PartialEq, Eq, Hash)]
 /// pub enum PyDataType {
 ///     #[pyo3(name = "FLOAT")]
@@ -52,18 +48,15 @@ pub fn gen_stub_pyclass_enum(_attr: TokenStream, item: TokenStream) -> TokenStre
 /// Embed metadata for Python stub file generation for `#[pymethods]` macro
 ///
 /// ```
-/// # use pyo3_stub_gen_derive::*;
-/// # use pyo3::*;
-/// # #[gen_stub_pyclass]
-/// # #[pyclass]
-/// # struct PyAddOp {}
-/// # #[pyclass]
-/// # struct Expression {}
-/// #[gen_stub_pymethods]
-/// #[pymethods]
-/// impl PyAddOp {
+/// #[pyo3_stub_gen_derive::gen_stub_pyclass]
+/// #[pyo3::pyclass]
+/// struct A {}
+///
+/// #[pyo3_stub_gen_derive::gen_stub_pymethods]
+/// #[pyo3::pymethods]
+/// impl A {
 ///     #[getter]
-///     fn get_terms(&self) -> Vec<Expression> {
+///     fn f(&self) -> Vec<u32> {
 ///        todo!()
 ///     }
 /// }
@@ -78,15 +71,10 @@ pub fn gen_stub_pymethods(_attr: TokenStream, item: TokenStream) -> TokenStream
 /// Embed metadata for Python stub file generation for `#[pyfunction]` macro
 ///
 /// ```
-/// # use pyo3_stub_gen_derive::*;
-/// # use pyo3::*;
-/// # #[pyclass]
-/// # #[derive(Clone)]
-/// # pub struct Expression {}
-/// #[gen_stub_pyfunction]
-/// #[pyfunction]
-/// #[pyo3(name = "is_linear")]
-/// pub fn py_is_linear(expr: Expression) -> bool {
+/// #[pyo3_stub_gen_derive::gen_stub_pyfunction]
+/// #[pyo3::pyfunction]
+/// #[pyo3(name = "is_odd")]
+/// pub fn is_odd(x: u32) -> bool {
 ///     todo!()
 /// }
 /// ```
@@ -95,15 +83,10 @@ pub fn gen_stub_pymethods(_attr: TokenStream, item: TokenStream) -> TokenStream
 /// If you want to append this function to another module, add `module` attribute.
 ///
 /// ```
-/// # use pyo3_stub_gen_derive::*;
-/// # use pyo3::*;
-/// # #[pyclass]
-/// # #[derive(Clone)]
-/// # pub struct Expression {}
-/// #[gen_stub_pyfunction(module = "my_module.experimental")]
-/// #[pyfunction]
-/// #[pyo3(name = "is_linear")]
-/// pub fn py_is_linear(expr: Expression) -> bool {
+/// #[pyo3_stub_gen_derive::gen_stub_pyfunction(module = "my_module.experimental")]
+/// #[pyo3::pyfunction]
+/// #[pyo3(name = "is_odd")]
+/// pub fn is_odd(x: u32) -> bool {
 ///     todo!()
 /// }
 /// ```

diff --git a/pyo3-stub-gen-testing-mixed/python/pyo3_stub_gen_testing_mixed/my_rust_pkg.pyi b/pyo3-stub-gen-testing-mixed/python/pyo3_stub_gen_testing_mixed/my_rust_pkg.pyi
@@ -1,8 +1,6 @@
 # This file is automatically generated by pyo3_stub_gen
 # ruff: noqa: E501, F401
 
-from typing import final, Any, List, Dict, Sequence, Mapping
-from enum import Enum, auto
 
 def sum_as_string(a:int,b:int) -> str:
     r"""

diff --git a/pyo3-stub-gen-testing-pure/pyo3_stub_gen_testing_pure.pyi b/pyo3-stub-gen-testing-pure/pyo3_stub_gen_testing_pure.pyi
@@ -1,10 +1,15 @@
 # This file is automatically generated by pyo3_stub_gen
 # ruff: noqa: E501, F401
 
-from typing import final, Any, List, Dict, Sequence, Mapping
-from enum import Enum, auto
+import typing
 
-def sum_as_string(a:int,b:int) -> str:
+def create_dict(n:int) -> dict[int, list[int]]:
+    ...
+
+def read_dict(dict:typing.Mapping[int, typing.Mapping[int, int]]) -> None:
+    ...
+
+def sum(v:typing.Sequence[int]) -> int:
     r"""
     Returns the sum of two numbers as a string.
     """

diff --git a/pyo3-stub-gen-testing-pure/src/lib.rs b/pyo3-stub-gen-testing-pure/src/lib.rs
@@ -3,12 +3,33 @@ mod readme {}
 
 use pyo3::{exceptions::PyRuntimeError, prelude::*};
 use pyo3_stub_gen::{create_exception, define_stub_info_gatherer, derive::gen_stub_pyfunction};
+use std::collections::HashMap;
 
 /// Returns the sum of two numbers as a string.
 #[gen_stub_pyfunction]
 #[pyfunction]
-fn sum_as_string(a: usize, b: usize) -> PyResult<String> {
-    Ok((a + b).to_string())
+fn sum(v: Vec<u32>) -> u32 {
+    v.iter().sum()
+}
+
+#[gen_stub_pyfunction]
+#[pyfunction]
+fn read_dict(dict: HashMap<usize, HashMap<usize, usize>>) {
+    for (k, v) in dict {
+        for (k2, v2) in v {
+            println!("{} {} {}", k, k2, v2);
+        }
+    }
+}
+
+#[gen_stub_pyfunction]
+#[pyfunction]
+fn create_dict(n: usize) -> HashMap<usize, Vec<usize>> {
+    let mut dict = HashMap::new();
+    for i in 0..n {
+        dict.insert(i, (0..i).collect());
+    }
+    dict
 }
 
 create_exception!(pyo3_stub_gen_testing_pure, MyError, PyRuntimeError);
@@ -17,7 +38,9 @@ create_exception!(pyo3_stub_gen_testing_pure, MyError, PyRuntimeError);
 #[pymodule]
 fn pyo3_stub_gen_testing_pure(m: &Bound<PyModule>) -> PyResult<()> {
     m.add("MyError", m.py().get_type_bound::<MyError>())?;
-    m.add_function(wrap_pyfunction!(sum_as_string, m)?)?;
+    m.add_function(wrap_pyfunction!(sum, m)?)?;
+    m.add_function(wrap_pyfunction!(create_dict, m)?)?;
+    m.add_function(wrap_pyfunction!(read_dict, m)?)?;
     Ok(())
 }
 

diff --git a/pyo3-stub-gen-testing-pure/tests/test_python.py b/pyo3-stub-gen-testing-pure/tests/test_python.py
@@ -1,5 +1,31 @@
-import pyo3_stub_gen_testing_pure
+from pyo3_stub_gen_testing_pure import sum, create_dict, read_dict
+import pytest
 
 
-def test_sum_as_string():
-    assert pyo3_stub_gen_testing_pure.sum_as_string(1, 2) == "3"
+def test_sum():
+    assert sum([1, 2]) == 3
+    assert sum((1, 2)) == 3
+
+
+def test_create_dict():
+    assert create_dict(3) == {0: [], 1: [0], 2: [0, 1]}
+
+
+def test_read_dict():
+    read_dict(
+        {
+            0: {
+                0: 1,
+            },
+            1: {
+                0: 2,
+                1: 3,
+            },
+        }
+    )
+
+    with pytest.raises(TypeError) as e:
+        read_dict({0: 1})  # type: ignore
+    assert (
+        str(e.value) == "argument 'dict': 'int' object cannot be converted to 'PyDict'"
+    )
diff --git a/pyo3-stub-gen/Cargo.toml b/pyo3-stub-gen/Cargo.toml
@@ -13,10 +13,11 @@ anyhow.workspace = true
 inventory.workspace = true
 itertools.workspace = true
 log.workspace = true
+maplit.workspace = true
 pyo3.workspace = true
 serde.workspace = true
 toml.workspace = true
 
 [dependencies.pyo3-stub-gen-derive]
-version = "0.4.1"
+version = "0.5.0"
 path = "../pyo3-stub-gen-derive"