1use pyo3::prelude::*;
2use pyo3::types::{PyDict, PyList, PySet, PyTuple};
3
4use crate::communication::append_usize_vec;
5use crate::create_numpy_pyany_serde;
6use crate::{
7 common::{
8 detect_python_type, get_python_type_byte, retrieve_python_type, NumpyDtype, PythonType,
9 },
10 communication::{append_usize, retrieve_usize},
11 PyAnySerde,
12};
13
14use super::{
15 BoolSerde, BytesSerde, ComplexSerde, FloatSerde, IntSerde, NumpySerde, NumpySerdeConfig,
16 PickleSerde, StringSerde,
17};
18
19#[derive(Clone)]
20pub struct DynamicSerde {
21 pickle_serde: PickleSerde,
22 int_serde: IntSerde,
23 float_serde: FloatSerde,
24 complex_serde: ComplexSerde,
25 boolean_serde: BoolSerde,
26 string_serde: StringSerde,
27 bytes_serde: BytesSerde,
28 numpy_i8_serde: NumpySerde<i8>,
29 numpy_i16_serde: NumpySerde<i16>,
30 numpy_i32_serde: NumpySerde<i32>,
31 numpy_i64_serde: NumpySerde<i64>,
32 numpy_u8_serde: NumpySerde<u8>,
33 numpy_u16_serde: NumpySerde<u16>,
34 numpy_u32_serde: NumpySerde<u32>,
35 numpy_u64_serde: NumpySerde<u64>,
36 numpy_f32_serde: NumpySerde<f32>,
37 numpy_f64_serde: NumpySerde<f64>,
38}
39
40impl DynamicSerde {
41 pub fn new() -> PyResult<Self> {
42 let pickle_serde = PickleSerde::new()?;
43 let int_serde = IntSerde {};
44 let float_serde = FloatSerde {};
45 let complex_serde = ComplexSerde {};
46 let boolean_serde = BoolSerde {};
47 let string_serde = StringSerde {};
48 let bytes_serde = BytesSerde {};
49 let numpy_serde_config = NumpySerdeConfig::DYNAMIC {
50 preprocessor_fn: None,
51 postprocessor_fn: None,
52 };
53 let numpy_i8_serde = *create_numpy_pyany_serde!(i8, numpy_serde_config.clone());
54 let numpy_i16_serde = *create_numpy_pyany_serde!(i16, numpy_serde_config.clone());
55 let numpy_i32_serde = *create_numpy_pyany_serde!(i32, numpy_serde_config.clone());
56 let numpy_i64_serde = *create_numpy_pyany_serde!(i64, numpy_serde_config.clone());
57 let numpy_u8_serde = *create_numpy_pyany_serde!(u8, numpy_serde_config.clone());
58 let numpy_u16_serde = *create_numpy_pyany_serde!(u16, numpy_serde_config.clone());
59 let numpy_u32_serde = *create_numpy_pyany_serde!(u32, numpy_serde_config.clone());
60 let numpy_u64_serde = *create_numpy_pyany_serde!(u64, numpy_serde_config.clone());
61 let numpy_f32_serde = *create_numpy_pyany_serde!(f32, numpy_serde_config.clone());
62 let numpy_f64_serde = *create_numpy_pyany_serde!(f64, numpy_serde_config.clone());
63
64 Ok(DynamicSerde {
65 pickle_serde,
66 int_serde,
67 float_serde,
68 complex_serde,
69 boolean_serde,
70 string_serde,
71 bytes_serde,
72 numpy_i8_serde,
73 numpy_i16_serde,
74 numpy_i32_serde,
75 numpy_i64_serde,
76 numpy_u8_serde,
77 numpy_u16_serde,
78 numpy_u32_serde,
79 numpy_u64_serde,
80 numpy_f32_serde,
81 numpy_f64_serde,
82 })
83 }
84}
85
86impl PyAnySerde for DynamicSerde {
87 fn append<'py>(
88 &mut self,
89 buf: &mut [u8],
90 mut offset: usize,
91 obj: &Bound<'py, PyAny>,
92 ) -> PyResult<usize> {
93 let python_type = detect_python_type(obj)?;
94 buf[offset] = get_python_type_byte(&python_type);
95 match python_type {
96 PythonType::BOOL => {
97 offset = self.boolean_serde.append(buf, offset, obj)?;
98 }
99 PythonType::INT => {
100 offset = self.int_serde.append(buf, offset, obj)?;
101 }
102 PythonType::FLOAT => {
103 offset = self.float_serde.append(buf, offset, obj)?;
104 }
105 PythonType::COMPLEX => {
106 offset = self.complex_serde.append(buf, offset, obj)?;
107 }
108 PythonType::STRING => {
109 offset = self.string_serde.append(buf, offset, obj)?;
110 }
111 PythonType::BYTES => {
112 offset = self.bytes_serde.append(buf, offset, obj)?;
113 }
114 PythonType::NUMPY { dtype } => match dtype {
115 NumpyDtype::INT8 => {
116 offset = self.numpy_i8_serde.append(buf, offset, obj)?;
117 }
118 NumpyDtype::INT16 => {
119 offset = self.numpy_i16_serde.append(buf, offset, obj)?;
120 }
121 NumpyDtype::INT32 => {
122 offset = self.numpy_i32_serde.append(buf, offset, obj)?;
123 }
124 NumpyDtype::INT64 => {
125 offset = self.numpy_i64_serde.append(buf, offset, obj)?;
126 }
127 NumpyDtype::UINT8 => {
128 offset = self.numpy_u8_serde.append(buf, offset, obj)?;
129 }
130 NumpyDtype::UINT16 => {
131 offset = self.numpy_u16_serde.append(buf, offset, obj)?;
132 }
133 NumpyDtype::UINT32 => {
134 offset = self.numpy_u32_serde.append(buf, offset, obj)?;
135 }
136 NumpyDtype::UINT64 => {
137 offset = self.numpy_u64_serde.append(buf, offset, obj)?;
138 }
139 NumpyDtype::FLOAT32 => {
140 offset = self.numpy_f32_serde.append(buf, offset, obj)?;
141 }
142 NumpyDtype::FLOAT64 => {
143 offset = self.numpy_f64_serde.append(buf, offset, obj)?;
144 }
145 },
146 PythonType::LIST => {
147 let list = obj.downcast::<PyList>()?;
148 offset = append_usize(buf, offset, list.len());
149 for item in list.iter() {
150 offset = self.append(buf, offset, &item)?;
151 }
152 }
153 PythonType::SET => {
154 let set = obj.downcast::<PySet>()?;
155 offset = append_usize(buf, offset, set.len());
156 for item in set.iter() {
157 offset = self.append(buf, offset, &item)?;
158 }
159 }
160 PythonType::TUPLE => {
161 let tuple = obj.downcast::<PyTuple>()?;
162 offset = append_usize(buf, offset, tuple.len());
163 for item in tuple.iter() {
164 offset = self.append(buf, offset, &item)?;
165 }
166 }
167 PythonType::DICT => {
168 let dict = obj.downcast::<PyDict>()?;
169 offset = append_usize(buf, offset, dict.len());
170 for (key, value) in dict.iter() {
171 offset = self.append(buf, offset, &key)?;
172 offset = self.append(buf, offset, &value)?;
173 }
174 }
175 PythonType::OTHER => {
176 offset = self.pickle_serde.append(buf, offset, obj)?;
177 }
178 };
179 Ok(offset)
180 }
181
182 fn append_vec<'py>(
183 &mut self,
184 v: &mut Vec<u8>,
185 start_addr: Option<usize>,
186 obj: &Bound<'py, PyAny>,
187 ) -> PyResult<()> {
188 let python_type = detect_python_type(obj)?;
189 v.push(get_python_type_byte(&python_type));
190 match python_type {
191 PythonType::BOOL => {
192 self.boolean_serde.append_vec(v, start_addr, obj)?;
193 }
194 PythonType::INT => {
195 self.int_serde.append_vec(v, start_addr, obj)?;
196 }
197 PythonType::FLOAT => {
198 self.float_serde.append_vec(v, start_addr, obj)?;
199 }
200 PythonType::COMPLEX => {
201 self.complex_serde.append_vec(v, start_addr, obj)?;
202 }
203 PythonType::STRING => {
204 self.string_serde.append_vec(v, start_addr, obj)?;
205 }
206 PythonType::BYTES => {
207 self.bytes_serde.append_vec(v, start_addr, obj)?;
208 }
209 PythonType::NUMPY { dtype } => match dtype {
210 NumpyDtype::INT8 => {
211 self.numpy_i8_serde.append_vec(v, start_addr, obj)?;
212 }
213 NumpyDtype::INT16 => {
214 self.numpy_i16_serde.append_vec(v, start_addr, obj)?;
215 }
216 NumpyDtype::INT32 => {
217 self.numpy_i32_serde.append_vec(v, start_addr, obj)?;
218 }
219 NumpyDtype::INT64 => {
220 self.numpy_i64_serde.append_vec(v, start_addr, obj)?;
221 }
222 NumpyDtype::UINT8 => {
223 self.numpy_u8_serde.append_vec(v, start_addr, obj)?;
224 }
225 NumpyDtype::UINT16 => {
226 self.numpy_u16_serde.append_vec(v, start_addr, obj)?;
227 }
228 NumpyDtype::UINT32 => {
229 self.numpy_u32_serde.append_vec(v, start_addr, obj)?;
230 }
231 NumpyDtype::UINT64 => {
232 self.numpy_u64_serde.append_vec(v, start_addr, obj)?;
233 }
234 NumpyDtype::FLOAT32 => {
235 self.numpy_f32_serde.append_vec(v, start_addr, obj)?;
236 }
237 NumpyDtype::FLOAT64 => {
238 self.numpy_f64_serde.append_vec(v, start_addr, obj)?;
239 }
240 },
241 PythonType::LIST => {
242 let list = obj.downcast::<PyList>()?;
243 append_usize_vec(v, list.len());
244 for item in list.iter() {
245 self.append_vec(v, start_addr, &item)?;
246 }
247 }
248 PythonType::SET => {
249 let set = obj.downcast::<PyList>()?;
250 append_usize_vec(v, set.len());
251 for item in set.iter() {
252 self.append_vec(v, start_addr, &item)?;
253 }
254 }
255 PythonType::TUPLE => {
256 let tuple = obj.downcast::<PyList>()?;
257 append_usize_vec(v, tuple.len());
258 for item in tuple.iter() {
259 self.append_vec(v, start_addr, &item)?;
260 }
261 }
262 PythonType::DICT => {
263 let dict = obj.downcast::<PyDict>()?;
264 append_usize_vec(v, dict.len());
265 for (key, value) in dict.iter() {
266 self.append_vec(v, start_addr, &key)?;
267 self.append_vec(v, start_addr, &value)?;
268 }
269 }
270 PythonType::OTHER => {
271 self.pickle_serde.append_vec(v, start_addr, obj)?;
272 }
273 };
274 Ok(())
275 }
276
277 fn retrieve<'py>(
278 &mut self,
279 py: Python<'py>,
280 buf: &[u8],
281 offset: usize,
282 ) -> PyResult<(Bound<'py, PyAny>, usize)> {
283 let (python_type, mut offset) = retrieve_python_type(buf, offset)?;
284 let obj;
285 match python_type {
286 PythonType::BOOL => {
287 (obj, offset) = self.boolean_serde.retrieve(py, buf, offset)?;
288 }
289 PythonType::INT => {
290 (obj, offset) = self.int_serde.retrieve(py, buf, offset)?;
291 }
292 PythonType::FLOAT => {
293 (obj, offset) = self.float_serde.retrieve(py, buf, offset)?;
294 }
295 PythonType::COMPLEX => {
296 (obj, offset) = self.complex_serde.retrieve(py, buf, offset)?;
297 }
298 PythonType::STRING => {
299 (obj, offset) = self.string_serde.retrieve(py, buf, offset)?;
300 }
301 PythonType::BYTES => {
302 (obj, offset) = self.bytes_serde.retrieve(py, buf, offset)?;
303 }
304 PythonType::NUMPY { dtype } => match dtype {
305 NumpyDtype::INT8 => {
306 (obj, offset) = self.numpy_i8_serde.retrieve(py, buf, offset)?;
307 }
308 NumpyDtype::INT16 => {
309 (obj, offset) = self.numpy_i16_serde.retrieve(py, buf, offset)?;
310 }
311 NumpyDtype::INT32 => {
312 (obj, offset) = self.numpy_i32_serde.retrieve(py, buf, offset)?;
313 }
314 NumpyDtype::INT64 => {
315 (obj, offset) = self.numpy_i64_serde.retrieve(py, buf, offset)?;
316 }
317 NumpyDtype::UINT8 => {
318 (obj, offset) = self.numpy_u8_serde.retrieve(py, buf, offset)?;
319 }
320 NumpyDtype::UINT16 => {
321 (obj, offset) = self.numpy_u16_serde.retrieve(py, buf, offset)?;
322 }
323 NumpyDtype::UINT32 => {
324 (obj, offset) = self.numpy_u32_serde.retrieve(py, buf, offset)?;
325 }
326 NumpyDtype::UINT64 => {
327 (obj, offset) = self.numpy_u64_serde.retrieve(py, buf, offset)?;
328 }
329 NumpyDtype::FLOAT32 => {
330 (obj, offset) = self.numpy_f32_serde.retrieve(py, buf, offset)?;
331 }
332 NumpyDtype::FLOAT64 => {
333 (obj, offset) = self.numpy_f64_serde.retrieve(py, buf, offset)?;
334 }
335 },
336 PythonType::LIST => {
337 let list = PyList::empty(py);
338 let n_items;
339 (n_items, offset) = retrieve_usize(buf, offset)?;
340 for _ in 0..n_items {
341 let item;
342 (item, offset) = self.retrieve(py, buf, offset)?;
343 list.append(item)?;
344 }
345 obj = list.into_any();
346 }
347 PythonType::SET => {
348 let set = PySet::empty(py)?;
349 let n_items;
350 (n_items, offset) = retrieve_usize(buf, offset)?;
351 for _ in 0..n_items {
352 let item;
353 (item, offset) = self.retrieve(py, buf, offset)?;
354 set.add(item)?;
355 }
356 obj = set.into_any();
357 }
358 PythonType::TUPLE => {
359 let n_items;
360 (n_items, offset) = retrieve_usize(buf, offset)?;
361 let mut tuple_vec = Vec::with_capacity(n_items);
362 for _ in 0..n_items {
363 let item;
364 (item, offset) = self.retrieve(py, buf, offset)?;
365 tuple_vec.push(item);
366 }
367 obj = PyTuple::new(py, tuple_vec)?.into_any();
368 }
369 PythonType::DICT => {
370 let dict = PyDict::new(py);
371 let n_items;
372 (n_items, offset) = retrieve_usize(buf, offset)?;
373 for _ in 0..n_items {
374 let key;
375 (key, offset) = self.retrieve(py, buf, offset)?;
376 let value;
377 (value, offset) = self.retrieve(py, buf, offset)?;
378 dict.set_item(key, value)?;
379 }
380 obj = dict.into_any();
381 }
382 PythonType::OTHER => {
383 (obj, offset) = self.pickle_serde.retrieve(py, buf, offset)?;
384 }
385 };
386 Ok((obj, offset))
387 }
388}