1use pyo3::prelude::*;
2use pyo3::types::PyString;
3
4use dyn_clone::{clone_trait_object, DynClone};
5
6use crate::communication::{append_bool, append_bool_vec, retrieve_bool};
7use crate::pyany_serde_impl::{
8 get_numpy_serde, BoolSerde, BytesSerde, ComplexSerde, DataclassSerde, DictSerde, DynamicSerde,
9 FloatSerde, IntSerde, ListSerde, OptionSerde, PickleSerde, PythonSerdeSerde, SetSerde,
10 StringSerde, TupleSerde, TypedDictSerde, UnionSerde,
11};
12use crate::pyany_serde_type::PyAnySerdeType;
13use crate::PickleablePyAnySerdeType;
14
15pub trait PyAnySerde: DynClone {
16 fn append<'py>(
17 &mut self,
18 buf: &mut [u8],
19 offset: usize,
20 obj: &Bound<'py, PyAny>,
21 ) -> PyResult<usize>;
22 fn append_vec<'py>(
23 &mut self,
24 v: &mut Vec<u8>,
25 start_addr: Option<usize>,
26 obj: &Bound<'py, PyAny>,
27 ) -> PyResult<()>;
28 fn retrieve<'py>(
29 &mut self,
30 py: Python<'py>,
31 buf: &[u8],
32 offset: usize,
33 ) -> PyResult<(Bound<'py, PyAny>, usize)>;
34 fn append_option<'py>(
35 &mut self,
36 buf: &mut [u8],
37 mut offset: usize,
38 obj_option: &Option<&Bound<'py, PyAny>>,
39 ) -> PyResult<usize> {
40 if let Some(obj) = obj_option {
41 offset = append_bool(buf, offset, true);
42 offset = self.append(buf, offset, obj)?;
43 } else {
44 offset = append_bool(buf, offset, false);
45 }
46 Ok(offset)
47 }
48 fn append_option_vec<'py>(
49 &mut self,
50 v: &mut Vec<u8>,
51 start_addr: Option<usize>,
52 obj_option: &Option<&Bound<'py, PyAny>>,
53 ) -> PyResult<()> {
54 if let Some(obj) = obj_option {
55 append_bool_vec(v, true);
56 self.append_vec(v, start_addr, obj)?;
57 } else {
58 append_bool_vec(v, false);
59 }
60 Ok(())
61 }
62 fn retrieve_option<'py>(
63 &mut self,
64 py: Python<'py>,
65 buf: &[u8],
66 mut offset: usize,
67 ) -> PyResult<(Option<Bound<'py, PyAny>>, usize)> {
68 let is_some;
69 (is_some, offset) = retrieve_bool(buf, offset)?;
70 if is_some {
71 let obj;
72 (obj, offset) = self.retrieve(py, buf, offset)?;
73 Ok((Some(obj), offset))
74 } else {
75 Ok((None, offset))
76 }
77 }
78}
79
80clone_trait_object!(PyAnySerde);
81
82impl<'py, 'a> TryFrom<&'a Bound<'py, PyAnySerdeType>> for Box<dyn PyAnySerde> {
83 type Error = PyErr;
84
85 fn try_from(value: &'a Bound<'py, PyAnySerdeType>) -> Result<Self, Self::Error> {
86 value.as_any().extract::<PyAnySerdeType>()?.try_into()
87 }
88}
89
90impl<'a> TryFrom<&'a Py<PyAnySerdeType>> for Box<dyn PyAnySerde> {
91 type Error = PyErr;
92
93 fn try_from(value: &'a Py<PyAnySerdeType>) -> Result<Self, Self::Error> {
94 Python::with_gil(|py| value.extract::<PyAnySerdeType>(py)?.try_into())
95 }
96}
97
98impl TryFrom<PyAnySerdeType> for Box<dyn PyAnySerde> {
99 type Error = PyErr;
100
101 fn try_from(value: PyAnySerdeType) -> Result<Self, Self::Error> {
102 (&value).try_into()
103 }
104}
105
106impl<'a> TryFrom<&'a PyAnySerdeType> for Box<dyn PyAnySerde> {
107 type Error = PyErr;
108
109 fn try_from(value: &'a PyAnySerdeType) -> Result<Self, Self::Error> {
110 Ok(match value {
111 PyAnySerdeType::BOOL {} => Box::new(BoolSerde {}),
112 PyAnySerdeType::BYTES {} => Box::new(BytesSerde {}),
113 PyAnySerdeType::COMPLEX {} => Box::new(ComplexSerde {}),
114 PyAnySerdeType::DATACLASS {
115 clazz,
116 init_strategy,
117 field_serde_type_dict,
118 } => Python::with_gil::<_, PyResult<_>>(|py| {
119 Ok(Box::new(DataclassSerde::new(
120 clazz.clone_ref(py),
121 init_strategy.clone(),
122 field_serde_type_dict
123 .iter()
124 .map(|(field, field_serde_type)| {
125 field_serde_type.try_into().map(|pyany_serde| {
126 (PyString::new(py, field.as_str()).unbind(), pyany_serde)
127 })
128 })
129 .collect::<PyResult<_>>()?,
130 )?))
131 })?,
132 PyAnySerdeType::DICT {
133 keys_serde_type,
134 values_serde_type,
135 } => Python::with_gil::<_, PyResult<_>>(|py| {
136 Ok(Box::new(DictSerde {
137 keys_serde: keys_serde_type.bind(py).try_into()?,
138 values_serde: values_serde_type.bind(py).try_into()?,
139 }))
140 })?,
141 PyAnySerdeType::DYNAMIC {} => Box::new(DynamicSerde::new()?),
142 PyAnySerdeType::FLOAT {} => Box::new(FloatSerde {}),
143 PyAnySerdeType::INT {} => Box::new(IntSerde {}),
144 PyAnySerdeType::LIST { items_serde_type } => Box::new(ListSerde {
145 items_serde: items_serde_type.try_into()?,
146 }),
147 PyAnySerdeType::NUMPY { dtype, config } => {
148 get_numpy_serde(dtype.clone(), config.clone())
149 }
150
151 PyAnySerdeType::OPTION { value_serde_type } => Box::new(OptionSerde {
152 value_serde: value_serde_type.try_into()?,
153 }),
154 PyAnySerdeType::PICKLE {} => Box::new(PickleSerde::new()?),
155 PyAnySerdeType::PYTHONSERDE { python_serde } => {
156 Python::with_gil::<_, PyResult<_>>(|py| {
157 Ok(Box::new(PythonSerdeSerde {
158 python_serde: python_serde.clone_ref(py),
159 }))
160 })?
161 }
162 PyAnySerdeType::SET { items_serde_type } => Box::new(SetSerde {
163 items_serde: items_serde_type.try_into()?,
164 }),
165 PyAnySerdeType::STRING {} => Box::new(StringSerde {}),
166 PyAnySerdeType::TUPLE { item_serde_types } => Box::new(TupleSerde {
167 item_serdes: item_serde_types
168 .into_iter()
169 .map(|item| item.try_into())
170 .collect::<PyResult<_>>()?,
171 }),
172 PyAnySerdeType::TYPEDDICT {
173 key_serde_type_dict,
174 } => Python::with_gil::<_, PyResult<_>>(|py| {
175 let serde_kv_list = key_serde_type_dict
176 .into_iter()
177 .map(|(key, item_serde_type)| {
178 item_serde_type.try_into().map(|pyany_serde| {
179 (PyString::new(py, key.as_str()).unbind(), pyany_serde)
180 })
181 })
182 .collect::<PyResult<_>>()?;
183 Ok(Box::new(TypedDictSerde { serde_kv_list }) as Box<dyn PyAnySerde>)
184 })?,
185 PyAnySerdeType::UNION {
186 option_serde_types,
187 option_choice_fn,
188 } => Python::with_gil::<_, PyResult<_>>(|py| {
189 Ok(Box::new(UnionSerde {
190 option_serdes: option_serde_types
191 .into_iter()
192 .map(|item| item.try_into())
193 .collect::<PyResult<_>>()?,
194 option_choice_fn: option_choice_fn.clone_ref(py),
195 }))
196 })?,
197 })
198 }
199}
200
201impl<'py> FromPyObject<'py> for Box<dyn PyAnySerde> {
202 fn extract_bound(ob: &Bound<'py, PyAny>) -> PyResult<Self> {
203 ob.extract::<PyAnySerdeType>()
204 .or_else(|_| {
205 ob.extract::<PickleablePyAnySerdeType>()
206 .map(|v| v.0.unwrap().unwrap())
207 })?
208 .try_into()
209 }
210}
211
212pub enum DynPyAnySerdeOption {
213 Some(Box<dyn PyAnySerde>),
214 None,
215}
216
217impl From<DynPyAnySerdeOption> for Option<Box<dyn PyAnySerde>> {
218 fn from(value: DynPyAnySerdeOption) -> Self {
219 match value {
220 DynPyAnySerdeOption::Some(pyany_serde) => Some(pyany_serde),
221 DynPyAnySerdeOption::None => None,
222 }
223 }
224}
225
226impl<'py> FromPyObject<'py> for DynPyAnySerdeOption {
227 fn extract_bound(ob: &Bound<'py, PyAny>) -> PyResult<Self> {
228 ob.extract::<Option<PyAnySerdeType>>()
229 .or_else(|_| {
230 ob.extract::<PickleablePyAnySerdeType>()
231 .map(|v| v.0.unwrap())
232 })
233 .map(|pyany_serde_type_option| {
234 pyany_serde_type_option
235 .map(|pyany_serde_type| {
236 pyany_serde_type
237 .try_into()
238 .map(|pyany_serde| DynPyAnySerdeOption::Some(pyany_serde))
239 })
240 .unwrap_or(Ok(DynPyAnySerdeOption::None))
241 })?
242 }
243}