monarch_types/
pyobject.rs

1/*
2 * Copyright (c) Meta Platforms, Inc. and affiliates.
3 * All rights reserved.
4 *
5 * This source code is licensed under the BSD-style license found in the
6 * LICENSE file in the root directory of this source tree.
7 */
8
9use pyo3::prelude::*;
10use pyo3::types::PyBytes;
11use serde::Deserialize;
12use serde::Serialize;
13use typeuri::Named;
14
15#[derive(Debug, Clone, Serialize, Deserialize, Named)]
16pub struct PickledPyObject {
17    #[serde(with = "serde_bytes")]
18    bytes: Vec<u8>,
19    cloudpickle: bool,
20}
21
22impl PickledPyObject {
23    pub fn pickle(obj: &Bound<'_, PyAny>) -> PyResult<Self> {
24        Self::pickle_impl(obj, false)
25    }
26
27    pub fn cloudpickle(obj: &Bound<'_, PyAny>) -> PyResult<Self> {
28        Self::pickle_impl(obj, true)
29    }
30
31    fn module(cloudpickle: bool) -> &'static str {
32        if cloudpickle { "cloudpickle" } else { "pickle" }
33    }
34
35    fn pickle_impl(obj: &Bound<'_, PyAny>, cloudpickle: bool) -> PyResult<Self> {
36        let module = Self::module(cloudpickle);
37        let bytes = obj
38            .py()
39            .import(module)?
40            .call_method1("dumps", (obj,))?
41            .cast_into::<PyBytes>()?
42            .as_bytes()
43            .to_vec();
44        Ok(Self { bytes, cloudpickle })
45    }
46
47    pub fn unpickle<'py>(&self, py: Python<'py>) -> PyResult<Bound<'py, PyAny>> {
48        py.import(Self::module(self.cloudpickle))?
49            .call_method1("loads", (self.bytes.as_slice(),))
50    }
51}
52
53impl TryFrom<&Bound<'_, PyAny>> for PickledPyObject {
54    type Error = PyErr;
55    fn try_from(obj: &Bound<'_, PyAny>) -> PyResult<Self> {
56        Self::pickle(obj)
57    }
58}
59
60impl TryFrom<Bound<'_, PyAny>> for PickledPyObject {
61    type Error = PyErr;
62    fn try_from(obj: Bound<'_, PyAny>) -> PyResult<Self> {
63        Self::pickle(&obj)
64    }
65}
66
67impl FromPyObject<'_, '_> for PickledPyObject {
68    type Error = PyErr;
69
70    fn extract(obj: pyo3::Borrowed<'_, '_, PyAny>) -> PyResult<Self> {
71        PickledPyObject::pickle(&obj)
72    }
73}
74
75impl<'py> IntoPyObject<'py> for &PickledPyObject {
76    type Target = PyAny;
77    type Output = Bound<'py, Self::Target>;
78    type Error = PyErr;
79
80    fn into_pyobject(self, py: Python<'py>) -> Result<Self::Output, Self::Error> {
81        self.unpickle(py)
82    }
83}
monarch_types/pyobject.rs

monarch_types/
pyobject.rs
