pyany_serde/pyany_serde_impl/
union_serde.rs

1use 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}