Make YRoomClientOptions picklable

stefanw · Sep 27, 2023 · 3044cfc · 3044cfc
1 parent d009485
commit 3044cfc
Show file tree

Hide file tree

Showing 2 changed files with 37 additions and 7 deletions.
diff --git a/src/roomsync.rs b/src/roomsync.rs
@@ -5,7 +5,7 @@ use std::{
 
 use pyo3::{
     prelude::*,
-    types::{PyBytes, PyDict, PyList},
+    types::{PyBytes, PyDict, PyList, PyTuple},
 };
 
 use lib0::{
@@ -230,7 +230,7 @@ impl FromPyObject<'_> for YRoomSettings {
 }
 
 #[derive(Clone, Debug)]
-#[pyclass]
+#[pyclass(module = "yroom")]
 pub struct YRoomClientOptions {
     allow_write: bool,
     allow_write_awareness: bool,
@@ -255,6 +255,28 @@ impl YRoomClientOptions {
             allow_write_awareness,
         }
     }
+
+    #[getter]
+    pub fn allow_write(&self) -> bool {
+        self.allow_write
+    }
+    #[getter]
+    pub fn allow_write_awareness(&self) -> bool {
+        self.allow_write_awareness
+    }
+
+    pub fn __getstate__(&self) -> PyResult<PyObject> {
+        Python::with_gil(|py| Ok((self.allow_write, self.allow_write_awareness).to_object(py)))
+    }
+
+    fn __setstate__(&mut self, state_tuple: PyObject) -> PyResult<()> {
+        Python::with_gil(|py| {
+            let tuple: &PyTuple = state_tuple.extract(py)?;
+            self.allow_write = tuple.get_item(0)?.extract::<bool>()?;
+            self.allow_write_awareness = tuple.get_item(1)?.extract::<bool>()?;
+            Ok(())
+        })
+    }
 }
 
 #[pyclass]
@@ -583,11 +605,11 @@ impl YRoom {
 
     fn connect(&mut self, conn_id: u64, options: YRoomClientOptions) -> YRoomMessage {
         {
-        let connections = self.connections.lock();
-        connections
-            .unwrap()
-            .entry(conn_id)
-            .or_insert_with(HashSet::new);
+            let connections = self.connections.lock();
+            connections
+                .unwrap()
+                .entry(conn_id)
+                .or_insert_with(HashSet::new);
         }
         let mut encoder = EncoderWrapper::new(
             &self.settings.protocol_version,

diff --git a/tests/test_yroom.py b/tests/test_yroom.py
@@ -1,4 +1,5 @@
 import json
+from copy import deepcopy
 
 import y_py as Y
 
@@ -396,3 +397,10 @@ def make_awareness_update(clients):
     message = manager.handle_message(room_name, client_3, payload_3, options)
     assert message.payloads == []
     assert message.broadcast_payloads == []
+
+
+def test_pickle_client_options():
+    options = YRoomClientOptions(allow_write=False, allow_write_awareness=False)
+    copied_options = deepcopy(options)
+    assert copied_options.allow_write is False
+    assert copied_options.allow_write_awareness is False