pyo3_stub_gen/
stub_type.rs

1mod builtins;
2mod collections;
3mod pyo3;
4
5#[cfg(feature = "numpy")]
6mod numpy;
7
8#[cfg(feature = "either")]
9mod either;
10
11use maplit::hashset;
12use std::{collections::HashSet, fmt, ops};
13
14#[derive(Debug, Clone, PartialEq, Eq, PartialOrd, Ord, Default, Hash)]
15pub enum ModuleRef {
16    Named(String),
17
18    /// Default module that PyO3 creates.
19    ///
20    /// - For pure Rust project, the default module name is the crate name specified in `Cargo.toml`
21    ///   or `project.name` specified in `pyproject.toml`
22    /// - For mixed Rust/Python project, the default module name is `tool.maturin.module-name` specified in `pyproject.toml`
23    ///
24    /// Because the default module name cannot be known at compile time, it will be resolved at the time of the stub file generation.
25    /// This is a placeholder for the default module name.
26    #[default]
27    Default,
28}
29
30impl ModuleRef {
31    pub fn get(&self) -> Option<&str> {
32        match self {
33            Self::Named(name) => Some(name),
34            Self::Default => None,
35        }
36    }
37}
38
39impl From<&str> for ModuleRef {
40    fn from(s: &str) -> Self {
41        Self::Named(s.to_string())
42    }
43}
44
45/// Type information for creating Python stub files annotated by [PyStubType] trait.
46#[derive(Debug, Clone, PartialEq, Eq)]
47pub struct TypeInfo {
48    /// The Python type name.
49    pub name: String,
50
51    /// Python modules must be imported in the stub file.
52    ///
53    /// For example, when `name` is `typing.Sequence[int]`, `import` should contain `typing`.
54    /// This makes it possible to use user-defined types in the stub file.
55    pub import: HashSet<ModuleRef>,
56}
57
58impl fmt::Display for TypeInfo {
59    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
60        write!(f, "{}", self.name)
61    }
62}
63
64impl TypeInfo {
65    /// A `None` type annotation.
66    pub fn none() -> Self {
67        // NOTE: since 3.10, NoneType is provided from types module,
68        // but there is no corresponding definitions prior to 3.10.
69        Self {
70            name: "None".to_string(),
71            import: HashSet::new(),
72        }
73    }
74
75    /// A `typing.Any` type annotation.
76    pub fn any() -> Self {
77        Self {
78            name: "typing.Any".to_string(),
79            import: hashset! { "typing".into() },
80        }
81    }
82
83    /// A `list[Type]` type annotation.
84    pub fn list_of<T: PyStubType>() -> Self {
85        let TypeInfo { name, mut import } = T::type_output();
86        import.insert("builtins".into());
87        TypeInfo {
88            name: format!("builtins.list[{}]", name),
89            import,
90        }
91    }
92
93    /// A `set[Type]` type annotation.
94    pub fn set_of<T: PyStubType>() -> Self {
95        let TypeInfo { name, mut import } = T::type_output();
96        import.insert("builtins".into());
97        TypeInfo {
98            name: format!("builtins.set[{}]", name),
99            import,
100        }
101    }
102
103    /// A `dict[Type]` type annotation.
104    pub fn dict_of<K: PyStubType, V: PyStubType>() -> Self {
105        let TypeInfo {
106            name: name_k,
107            mut import,
108        } = K::type_output();
109        let TypeInfo {
110            name: name_v,
111            import: import_v,
112        } = V::type_output();
113        import.extend(import_v);
114        import.insert("builtins".into());
115        TypeInfo {
116            name: format!("builtins.set[{}, {}]", name_k, name_v),
117            import,
118        }
119    }
120
121    /// A type annotation of a built-in type provided from `builtins` module, such as `int`, `str`, or `float`. Generic builtin types are also possible, such as `dict[str, str]`.
122    pub fn builtin(name: &str) -> Self {
123        Self {
124            name: format!("builtins.{name}"),
125            import: hashset! { "builtins".into() },
126        }
127    }
128
129    /// Unqualified type.
130    pub fn unqualified(name: &str) -> Self {
131        Self {
132            name: name.to_string(),
133            import: hashset! {},
134        }
135    }
136
137    /// A type annotation of a type that must be imported. The type name must be qualified with the module name:
138    ///
139    /// ```
140    /// pyo3_stub_gen::TypeInfo::with_module("pathlib.Path", "pathlib".into());
141    /// ```
142    pub fn with_module(name: &str, module: ModuleRef) -> Self {
143        let mut import = HashSet::new();
144        import.insert(module);
145        Self {
146            name: name.to_string(),
147            import,
148        }
149    }
150}
151
152impl ops::BitOr for TypeInfo {
153    type Output = Self;
154
155    fn bitor(mut self, rhs: Self) -> Self {
156        self.import.extend(rhs.import);
157        Self {
158            name: format!("{} | {}", self.name, rhs.name),
159            import: self.import,
160        }
161    }
162}
163
164/// Implement [PyStubType]
165///
166/// ```rust
167/// use pyo3::*;
168/// use pyo3_stub_gen::{impl_stub_type, derive::*};
169///
170/// #[gen_stub_pyclass]
171/// #[pyclass]
172/// struct A;
173///
174/// #[gen_stub_pyclass]
175/// #[pyclass]
176/// struct B;
177///
178/// enum E {
179///     A(A),
180///     B(B),
181/// }
182/// impl_stub_type!(E = A | B);
183///
184/// struct X(A);
185/// impl_stub_type!(X = A);
186///
187/// struct Y {
188///    a: A,
189///    b: B,
190/// }
191/// impl_stub_type!(Y = (A, B));
192/// ```
193#[macro_export]
194macro_rules! impl_stub_type {
195    ($ty: ty = $($base:ty)|+) => {
196        impl ::pyo3_stub_gen::PyStubType for $ty {
197            fn type_output() -> ::pyo3_stub_gen::TypeInfo {
198                $(<$base>::type_output()) | *
199            }
200            fn type_input() -> ::pyo3_stub_gen::TypeInfo {
201                $(<$base>::type_input()) | *
202            }
203        }
204    };
205    ($ty:ty = $base:ty) => {
206        impl ::pyo3_stub_gen::PyStubType for $ty {
207            fn type_output() -> ::pyo3_stub_gen::TypeInfo {
208                <$base>::type_output()
209            }
210            fn type_input() -> ::pyo3_stub_gen::TypeInfo {
211                <$base>::type_input()
212            }
213        }
214    };
215}
216
217/// Annotate Rust types with Python type information.
218pub trait PyStubType {
219    /// The type to be used in the output signature, i.e. return type of the Python function or methods.
220    fn type_output() -> TypeInfo;
221
222    /// The type to be used in the input signature, i.e. the arguments of the Python function or methods.
223    ///
224    /// This defaults to the output type, but can be overridden for types that are not valid input types.
225    /// For example, `Vec::<T>::type_output` returns `list[T]` while `Vec::<T>::type_input` returns `typing.Sequence[T]`.
226    fn type_input() -> TypeInfo {
227        Self::type_output()
228    }
229}
230
231#[cfg(test)]
232mod test {
233    use super::*;
234    use maplit::hashset;
235    use std::collections::HashMap;
236    use test_case::test_case;
237
238    #[test_case(bool::type_input(), "builtins.bool", hashset! { "builtins".into() } ; "bool_input")]
239    #[test_case(<&str>::type_input(), "builtins.str", hashset! { "builtins".into() } ; "str_input")]
240    #[test_case(Vec::<u32>::type_input(), "typing.Sequence[builtins.int]", hashset! { "typing".into(), "builtins".into() } ; "Vec_u32_input")]
241    #[test_case(Vec::<u32>::type_output(), "builtins.list[builtins.int]", hashset! {  "builtins".into() } ; "Vec_u32_output")]
242    #[test_case(HashMap::<u32, String>::type_input(), "typing.Mapping[builtins.int, builtins.str]", hashset! { "typing".into(), "builtins".into() } ; "HashMap_u32_String_input")]
243    #[test_case(HashMap::<u32, String>::type_output(), "builtins.dict[builtins.int, builtins.str]", hashset! { "builtins".into() } ; "HashMap_u32_String_output")]
244    #[test_case(indexmap::IndexMap::<u32, String>::type_input(), "typing.Mapping[builtins.int, builtins.str]", hashset! { "typing".into(), "builtins".into() } ; "IndexMap_u32_String_input")]
245    #[test_case(indexmap::IndexMap::<u32, String>::type_output(), "builtins.dict[builtins.int, builtins.str]", hashset! { "builtins".into() } ; "IndexMap_u32_String_output")]
246    #[test_case(HashMap::<u32, Vec<u32>>::type_input(), "typing.Mapping[builtins.int, typing.Sequence[builtins.int]]", hashset! { "builtins".into(), "typing".into() } ; "HashMap_u32_Vec_u32_input")]
247    #[test_case(HashMap::<u32, Vec<u32>>::type_output(), "builtins.dict[builtins.int, builtins.list[builtins.int]]", hashset! { "builtins".into() } ; "HashMap_u32_Vec_u32_output")]
248    #[test_case(HashSet::<u32>::type_input(), "builtins.set[builtins.int]", hashset! { "builtins".into() } ; "HashSet_u32_input")]
249    #[test_case(indexmap::IndexSet::<u32>::type_input(), "builtins.set[builtins.int]", hashset! { "builtins".into() } ; "IndexSet_u32_input")]
250    fn test(tinfo: TypeInfo, name: &str, import: HashSet<ModuleRef>) {
251        assert_eq!(tinfo.name, name);
252        if import.is_empty() {
253            assert!(tinfo.import.is_empty());
254        } else {
255            assert_eq!(tinfo.import, import);
256        }
257    }
258}
pyo3_stub_gen/stub_type.rs

pyo3_stub_gen/
stub_type.rs
