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 hyperactor::Named;
10use pyo3::prelude::*;
11use pyo3::types::PyBytes;
12use serde::Deserialize;
13use serde::Serialize;
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<'py>(obj: &Bound<'py, PyAny>) -> PyResult<Self> {
24        Self::pickle_impl(obj, false)
25    }
26
27    pub fn cloudpickle<'py>(obj: &Bound<'py, 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<'py>(obj: &Bound<'py, 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            .downcast_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    fn extract_bound(obj: &Bound<'_, PyAny>) -> PyResult<Self> {
69        PickledPyObject::pickle(obj)
70    }
71}
72
73impl<'py> IntoPyObject<'py> for &PickledPyObject {
74    type Target = PyAny;
75    type Output = Bound<'py, Self::Target>;
76    type Error = PyErr;
77
78    fn into_pyobject(self, py: Python<'py>) -> Result<Self::Output, Self::Error> {
79        self.unpickle(py)
80    }
81}
monarch_types/pyobject.rs

monarch_types/
pyobject.rs
