pyany_serde/
communication.rs

1use std::mem::size_of;
2use std::os::raw::c_double;
3
4use pyo3::exceptions::asyncio::InvalidStateError;
5use pyo3::prelude::*;
6
7use paste::paste;
8
9use crate::PyAnySerde;
10
11macro_rules! define_primitive_communication {
12    ($type:ty) => {
13        paste! {
14            pub fn [<append_ $type>](buf: &mut [u8], offset: usize, val: $type) -> usize {
15                let end = offset + size_of::<$type>();
16                buf[offset..end].copy_from_slice(&val.to_ne_bytes());
17                end
18            }
19
20            pub fn [<append_ $type _vec>](v: &mut Vec<u8>, val: $type) {
21                v.extend_from_slice(&val.to_ne_bytes());
22            }
23
24            pub fn [<retrieve_ $type>](buf: &[u8], offset: usize) -> PyResult<($type, usize)> {
25                let end = offset + size_of::<$type>();
26                Ok(($type::from_ne_bytes(buf[offset..end].try_into()?), end))
27            }
28        }
29    };
30}
31
32define_primitive_communication!(usize);
33define_primitive_communication!(c_double);
34define_primitive_communication!(i64);
35define_primitive_communication!(u64);
36define_primitive_communication!(f32);
37define_primitive_communication!(f64);
38
39pub fn append_bool(buf: &mut [u8], offset: usize, val: bool) -> usize {
40    let end = offset + size_of::<u8>();
41    buf[offset..end].copy_from_slice(&(val as u8).to_ne_bytes());
42    end
43}
44
45pub fn append_bool_vec(v: &mut Vec<u8>, val: bool) {
46    v.extend_from_slice(&(val as u8).to_ne_bytes());
47}
48
49pub fn retrieve_bool(buf: &[u8], offset: usize) -> PyResult<(bool, usize)> {
50    let end = offset + size_of::<bool>();
51    let val = match u8::from_ne_bytes(buf[offset..end].try_into()?) {
52        0 => Ok(false),
53        1 => Ok(true),
54        v => Err(InvalidStateError::new_err(format!(
55            "tried to retrieve bool from shared_memory but got value {v}"
56        ))),
57    }?;
58    Ok((val, end))
59}
60
61pub fn append_bytes_vec(v: &mut Vec<u8>, bytes: &[u8]) {
62    append_usize_vec(v, bytes.len());
63    v.extend_from_slice(bytes);
64}
65
66pub fn append_string_vec(v: &mut Vec<u8>, s: &String) {
67    append_bytes_vec(v, s.as_bytes());
68}
69
70pub fn retrieve_string(buf: &[u8], offset: usize) -> PyResult<(String, usize)> {
71    let (string_bytes, offset) = retrieve_bytes(buf, offset)?;
72    Ok((String::from_utf8(string_bytes.to_vec())?, offset))
73}
74
75pub fn insert_bytes(buf: &mut [u8], offset: usize, bytes: &[u8]) -> usize {
76    let end = offset + bytes.len();
77    buf[offset..end].copy_from_slice(bytes);
78    end
79}
80
81pub fn append_bytes(buf: &mut [u8], offset: usize, bytes: &[u8]) -> usize {
82    let bytes_len = bytes.len();
83    let start = append_usize(buf, offset, bytes_len);
84    let end = start + bytes.len();
85    buf[start..end].copy_from_slice(bytes);
86    end
87}
88
89pub fn retrieve_bytes(buf: &[u8], offset: usize) -> PyResult<(&[u8], usize)> {
90    let (len, start) = retrieve_usize(buf, offset)?;
91    let end = start + len;
92    Ok((&buf[start..end], end))
93}
94
95pub fn append_python_option_bound<'py, F>(
96    buf: &mut [u8],
97    mut offset: usize,
98    obj_option: &Option<&Bound<'py, PyAny>>,
99    serde_option: &mut Option<&mut Box<dyn PyAnySerde>>,
100    err: F,
101) -> PyResult<usize>
102where
103    F: FnOnce() -> PyErr,
104{
105    if let Some(obj) = obj_option {
106        offset = append_bool(buf, offset, true);
107        offset = serde_option
108            .as_deref_mut()
109            .ok_or_else(err)?
110            .append(buf, offset, obj)?;
111    } else {
112        offset = append_bool(buf, offset, false);
113    }
114    Ok(offset)
115}
116
117pub fn append_python_option<'py, F>(
118    py: Python<'py>,
119    buf: &mut [u8],
120    mut offset: usize,
121    obj_option: &Option<&Py<PyAny>>,
122    serde_option: &mut Option<&mut Box<dyn PyAnySerde>>,
123    err: F,
124) -> PyResult<usize>
125where
126    F: FnOnce() -> PyErr,
127{
128    if let Some(obj) = obj_option {
129        offset = append_bool(buf, offset, true);
130        offset = serde_option
131            .as_deref_mut()
132            .ok_or_else(err)?
133            .append(buf, offset, obj.bind(py))?;
134    } else {
135        offset = append_bool(buf, offset, false);
136    }
137    Ok(offset)
138}
139
140pub fn retrieve_python_option<'py, F>(
141    py: Python<'py>,
142    buf: &mut [u8],
143    mut offset: usize,
144    serde_option: &mut Option<&mut Box<dyn PyAnySerde>>,
145    err: F,
146) -> PyResult<(Option<Bound<'py, PyAny>>, usize)>
147where
148    F: FnOnce() -> PyErr,
149{
150    let is_some;
151    (is_some, offset) = retrieve_bool(buf, offset)?;
152    let obj_option = if is_some {
153        let obj;
154        (obj, offset) = serde_option
155            .as_deref_mut()
156            .ok_or_else(err)?
157            .retrieve(py, buf, offset)?;
158        Some(obj)
159    } else {
160        None
161    };
162    Ok((obj_option, offset))
163}
pyany_serde/communication.rs

pyany_serde/
communication.rs
