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