pyany_serde/pyany_serde_impl/
pickle_serde.rs

1use pyo3::prelude::*;
2use pyo3::types::PyBytes;
3
4use crate::{
5    communication::{append_bytes, append_bytes_vec, retrieve_bytes},
6    PyAnySerde,
7};
8
9#[derive(Clone)]
10pub struct PickleSerde {
11    pickle_dumps: Py<PyAny>,
12    pickle_loads: Py<PyAny>,
13}
14
15impl PickleSerde {
16    pub fn new() -> PyResult<Self> {
17        Python::with_gil(|py| {
18            Ok(PickleSerde {
19                pickle_dumps: py.import("pickle")?.getattr("dumps")?.unbind(),
20                pickle_loads: py.import("pickle")?.getattr("loads")?.unbind(),
21            })
22        })
23    }
24}
25
26impl PyAnySerde for PickleSerde {
27    fn append<'py>(
28        &mut self,
29        buf: &mut [u8],
30        offset: usize,
31        obj: &Bound<'py, PyAny>,
32    ) -> PyResult<usize> {
33        Ok(append_bytes(
34            buf,
35            offset,
36            self.pickle_dumps
37                .bind(obj.py())
38                .call1((obj,))?
39                .downcast_into::<PyBytes>()?
40                .as_bytes(),
41        ))
42    }
43
44    fn append_vec<'py>(
45        &mut self,
46        v: &mut Vec<u8>,
47        _start_addr: Option<usize>,
48        obj: &Bound<'py, PyAny>,
49    ) -> PyResult<()> {
50        append_bytes_vec(
51            v,
52            self.pickle_dumps
53                .bind(obj.py())
54                .call1((obj,))?
55                .downcast_into::<PyBytes>()?
56                .as_bytes(),
57        );
58        Ok(())
59    }
60
61    fn retrieve<'py>(
62        &mut self,
63        py: Python<'py>,
64        buf: &[u8],
65        offset: usize,
66    ) -> PyResult<(Bound<'py, PyAny>, usize)> {
67        let (bytes, offset) = retrieve_bytes(buf, offset)?;
68        Ok((
69            self.pickle_loads
70                .bind(py)
71                .call1((PyBytes::new(py, bytes),))?,
72            offset,
73        ))
74    }
75}