pyo3_stub_gen/stub_type/
collections.rs

1use crate::stub_type::*;
2use std::collections::{BTreeMap, BTreeSet, HashMap};
3
4impl<T: PyStubType> PyStubType for Option<T> {
5    fn type_input() -> TypeInfo {
6        let TypeInfo { name, mut import } = T::type_input();
7        import.insert("typing".into());
8        TypeInfo {
9            name: format!("typing.Optional[{}]", name),
10            import,
11        }
12    }
13    fn type_output() -> TypeInfo {
14        let TypeInfo { name, mut import } = T::type_output();
15        import.insert("typing".into());
16        TypeInfo {
17            name: format!("typing.Optional[{}]", name),
18            import,
19        }
20    }
21}
22
23impl<T: PyStubType> PyStubType for Box<T> {
24    fn type_input() -> TypeInfo {
25        T::type_input()
26    }
27    fn type_output() -> TypeInfo {
28        T::type_output()
29    }
30}
31
32impl<T: PyStubType, E> PyStubType for Result<T, E> {
33    fn type_input() -> TypeInfo {
34        T::type_input()
35    }
36    fn type_output() -> TypeInfo {
37        T::type_output()
38    }
39}
40
41impl<T: PyStubType> PyStubType for Vec<T> {
42    fn type_input() -> TypeInfo {
43        let TypeInfo { name, mut import } = T::type_input();
44        import.insert("typing".into());
45        TypeInfo {
46            name: format!("typing.Sequence[{}]", name),
47            import,
48        }
49    }
50    fn type_output() -> TypeInfo {
51        TypeInfo::list_of::<T>()
52    }
53}
54
55impl<T: PyStubType, const N: usize> PyStubType for [T; N] {
56    fn type_input() -> TypeInfo {
57        let TypeInfo { name, mut import } = T::type_input();
58        import.insert("typing".into());
59        TypeInfo {
60            name: format!("typing.Sequence[{}]", name),
61            import,
62        }
63    }
64    fn type_output() -> TypeInfo {
65        TypeInfo::list_of::<T>()
66    }
67}
68
69impl<T: PyStubType, State> PyStubType for HashSet<T, State> {
70    fn type_output() -> TypeInfo {
71        TypeInfo::set_of::<T>()
72    }
73}
74
75impl<T: PyStubType> PyStubType for BTreeSet<T> {
76    fn type_output() -> TypeInfo {
77        TypeInfo::set_of::<T>()
78    }
79}
80
81impl<T: PyStubType> PyStubType for indexmap::IndexSet<T> {
82    fn type_output() -> TypeInfo {
83        TypeInfo::set_of::<T>()
84    }
85}
86
87macro_rules! impl_map_inner {
88    () => {
89        fn type_input() -> TypeInfo {
90            let TypeInfo {
91                name: key_name,
92                mut import,
93            } = Key::type_input();
94            let TypeInfo {
95                name: value_name,
96                import: value_import,
97            } = Value::type_input();
98            import.extend(value_import);
99            import.insert("typing".into());
100            TypeInfo {
101                name: format!("typing.Mapping[{}, {}]", key_name, value_name),
102                import,
103            }
104        }
105        fn type_output() -> TypeInfo {
106            let TypeInfo {
107                name: key_name,
108                mut import,
109            } = Key::type_output();
110            let TypeInfo {
111                name: value_name,
112                import: value_import,
113            } = Value::type_output();
114            import.extend(value_import);
115            import.insert("builtins".into());
116            TypeInfo {
117                name: format!("builtins.dict[{}, {}]", key_name, value_name),
118                import,
119            }
120        }
121    };
122}
123
124impl<Key: PyStubType, Value: PyStubType> PyStubType for BTreeMap<Key, Value> {
125    impl_map_inner!();
126}
127
128impl<Key: PyStubType, Value: PyStubType, State> PyStubType for HashMap<Key, Value, State> {
129    impl_map_inner!();
130}
131
132impl<Key: PyStubType, Value: PyStubType, State> PyStubType
133    for indexmap::IndexMap<Key, Value, State>
134{
135    impl_map_inner!();
136}
137
138macro_rules! impl_tuple {
139    ($($T:ident),*) => {
140        impl<$($T: PyStubType),*> PyStubType for ($($T),*) {
141            fn type_output() -> TypeInfo {
142                let mut merged = HashSet::new();
143                let mut names = Vec::new();
144                $(
145                let TypeInfo { name, import } = $T::type_output();
146                names.push(name);
147                merged.extend(import);
148                )*
149                TypeInfo {
150                    name: format!("tuple[{}]", names.join(", ")),
151                    import: merged,
152                }
153            }
154            fn type_input() -> TypeInfo {
155                let mut merged = HashSet::new();
156                let mut names = Vec::new();
157                $(
158                let TypeInfo { name, import } = $T::type_input();
159                names.push(name);
160                merged.extend(import);
161                )*
162                TypeInfo {
163                    name: format!("tuple[{}]", names.join(", ")),
164                    import: merged,
165                }
166            }
167        }
168    };
169}
170
171impl_tuple!(T1, T2);
172impl_tuple!(T1, T2, T3);
173impl_tuple!(T1, T2, T3, T4);
174impl_tuple!(T1, T2, T3, T4, T5);
175impl_tuple!(T1, T2, T3, T4, T5, T6);
176impl_tuple!(T1, T2, T3, T4, T5, T6, T7);
177impl_tuple!(T1, T2, T3, T4, T5, T6, T7, T8);
178impl_tuple!(T1, T2, T3, T4, T5, T6, T7, T8, T9);
pyo3_stub_gen/stub_type/collections.rs

pyo3_stub_gen/stub_type/
collections.rs
