pyany_serde/pyany_serde_impl/
union_serde.rs1use pyo3::exceptions::asyncio::InvalidStateError;
2use pyo3::prelude::*;
3use pyo3::types::PyFunction;
4
5use crate::{
6 communication::{append_usize, append_usize_vec, retrieve_usize},
7 PyAnySerde,
8};
9
10#[derive(Clone)]
11pub struct UnionSerde {
12 pub option_serdes: Vec<Box<dyn PyAnySerde>>,
13 pub option_choice_fn: Py<PyFunction>,
14}
15
16impl PyAnySerde for UnionSerde {
17 fn append<'py>(
18 &mut self,
19 buf: &mut [u8],
20 offset: usize,
21 obj: &Bound<'py, PyAny>,
22 ) -> PyResult<usize> {
23 let serde_idx = self
24 .option_choice_fn
25 .bind(obj.py())
26 .call1((obj,))?
27 .extract::<usize>()?;
28 let offset = append_usize(buf, offset, serde_idx);
29 let pyany_serde = self.option_serdes.get_mut(serde_idx).ok_or_else(|| {
30 InvalidStateError::new_err(format!(
31 "Serde choice function returned {} which is not a valid choice index",
32 serde_idx
33 ))
34 })?;
35 pyany_serde.append(buf, offset, obj)
36 }
37
38 fn append_vec<'py>(
39 &mut self,
40 v: &mut Vec<u8>,
41 start_addr: Option<usize>,
42 obj: &Bound<'py, PyAny>,
43 ) -> PyResult<()> {
44 let serde_idx = self
45 .option_choice_fn
46 .bind(obj.py())
47 .call1((obj,))?
48 .extract::<usize>()?;
49 append_usize_vec(v, serde_idx);
50 let pyany_serde = self.option_serdes.get_mut(serde_idx).ok_or_else(|| {
51 InvalidStateError::new_err(format!(
52 "Serde choice function returned {} which is not a valid choice index",
53 serde_idx
54 ))
55 })?;
56 pyany_serde.append_vec(v, start_addr, obj)
57 }
58
59 fn retrieve<'py>(
60 &mut self,
61 py: Python<'py>,
62 buf: &[u8],
63 offset: usize,
64 ) -> PyResult<(Bound<'py, PyAny>, usize)> {
65 let (serde_idx, offset) = retrieve_usize(buf, offset)?;
66 let pyany_serde = self.option_serdes.get_mut(serde_idx).ok_or_else(|| {
67 InvalidStateError::new_err(format!(
68 "Deserialized serde idx {} which is not a valid choice index",
69 serde_idx
70 ))
71 })?;
72 pyany_serde.retrieve(py, buf, offset)
73 }
74}