Pickling and unpickling objects: a working example

[ChristopherRabotin]

If you're looking to pickle and unpickle objects on the Python side, here's a method that works with PyO3 0.18+ (and possibly earlier versions too). I thought I would post this because I didn't find an easy way to handle this.

1. Add the pythonize package, cf. https://github.com/davidhewitt/pythonize/ and https://docs.rs/pythonize
2. Derive (or implement if you're brave) Serialize and Deseralize on your object
3. Add the method __getstate__ as follows to pickle your object

    #[cfg(feature = "python")]
    fn __getstate__(&self, py: Python) -> Result<PyObject, NyxError> {
        pythonize(py, &self).map_err(|e| NyxError::CustomError(e.to_string()))
    }

4. Add the method __setstate__ as follows to unpickle your object:

    #[cfg(feature = "python")]
    fn __setstate__(&mut self, state: &PyAny) -> Result<(), ConfigError> {
        *self = depythonize(state).map_err(|e| ConfigError::InvalidConfig(e.to_string()))?;
        Ok(())
    }

5. Make sure to define a #[new] function that can take no input arguments and still return a valid object. This could be done by implementing Default, or by doing something a bit more complicated like this:

    #[cfg(feature = "python")]
    #[new]
    fn py_new(
        tau: Option<Duration>,
        sigma: Option<f64>,
        steady_state: Option<f64>,
        bias: Option<f64>,
        epoch: Option<Epoch>,
    ) -> Result<Self, ConfigError> {
        if tau.is_none() && sigma.is_none() && bias.is_none() {
            // We're called from pickle, return a non initialized state
            return Ok(Self::ZERO);
        } else if tau.is_none() || sigma.is_none() || bias.is_none() {
            return Err(ConfigError::InvalidConfig(format!(
                "tau, sigma, and bias must be specified"
            )));
        }

        let tau = tau.unwrap();
        let sigma = sigma.unwrap();
        let steady_state = steady_state.unwrap();

        if tau <= Duration::ZERO {
            return Err(ConfigError::InvalidConfig(format!(
                "tau must be positive but got {tau}"
            )));
        }

        Ok(Self {
            tau,
            bias_sigma: sigma,
            steady_state_sigma: steady_state,
            bias,
            epoch,
        })
    }

Finally, the following works in Python:

def test_gauss_markov():
    from_dict = GaussMarkov.loads(gm_data)
    assert len(from_dict) == 2

    from_list = GaussMarkov.loads(list(gm_data.values()))
    assert len(from_list) == 2

    unique_range = GaussMarkov.loads(gm_data["range_noise_km"])[0]
    assert unique_range.tau == Unit.Day * 1.0

    # NOTE: We can only pickle via the `dumps` and `loads` functions.
    # This looses the stateful info.
    import pickle

    pkld = pickle.dumps(unique_range)
    unpkl = pickle.loads(pkld)
    assert unpkl.tau == Unit.Day * 1.0

The above code lives in Nyx, a high fidelity astrodynamics library for mission design and orbit determination of spacecraft. The time management stuff (Unit.Day * 1.0) is from hifitime.

[adamreichold]

I think quickwit-oss/tantivy-py#97 might also be a nice real-world example of the above which notably also includes frozen classes.