Skip to content

Commit abf98b7

Browse files
committed
improve serialization error handling
1 parent 42b8ef7 commit abf98b7

File tree

1 file changed

+14
-9
lines changed

1 file changed

+14
-9
lines changed

src/gridcounts.rs

Lines changed: 14 additions & 9 deletions
Original file line numberDiff line numberDiff line change
@@ -254,19 +254,24 @@ impl GridCounts {
254254
}
255255

256256
fn __setstate__(&mut self, state: Bound<'_, PyBytes>) -> PyResult<()> {
257-
let (counts, shape, resolution, n_threads) = deserialize(state.as_bytes()).unwrap();
258-
self.counts = counts;
259-
self.shape = shape;
260-
self.resolution = resolution;
261-
self.set_n_threads(Some(n_threads))?;
257+
match deserialize(state.as_bytes()) {
258+
Ok((counts, shape, resolution, n_threads)) => {
259+
self.counts = counts;
260+
self.shape = shape;
261+
self.resolution = resolution;
262+
self.set_n_threads(Some(n_threads))?;
262263

263-
Ok(())
264+
Ok(())
265+
}
266+
Err(e) => Err(PyValueError::new_err(e.to_string())),
267+
}
264268
}
265269

266270
fn __getstate__<'py>(&self, py: Python<'py>) -> PyResult<Bound<'py, PyBytes>> {
267-
let to_bytes = (&self.counts, self.shape, self.resolution, self.n_threads);
268-
269-
Ok(PyBytes::new_bound(py, &serialize(&to_bytes).unwrap()))
271+
match serialize(&(&self.counts, self.shape, self.resolution, self.n_threads)) {
272+
Ok(bytes) => Ok(PyBytes::new_bound(py, &bytes)),
273+
Err(e) => Err(PyRuntimeError::new_err(e.to_string())),
274+
}
270275
}
271276

272277
fn __getnewargs_ex__(

0 commit comments

Comments
 (0)