feat: Implement set_custom_formatters as a one unified gateway function

binding/python/examples/cli/raftify_cli/cli.py

@@ -5,14 +5,7 @@
 from typing import Optional
 from raftify import (
     cli_main,
-    set_confchange_context_deserializer,
-    set_confchangev2_context_deserializer,
-    set_entry_context_deserializer,
-    set_entry_data_deserializer,
-    set_message_context_deserializer,
-    set_snapshot_data_deserializer,
-    set_fsm_deserializer,
-    set_log_entry_deserializer,
+    set_custom_formatters,
 )
 
 
@@ -53,15 +46,16 @@ def register_custom_deserializer() -> None:
     Initialize the custom deserializers.
     """
 
-    set_confchange_context_deserializer(pickle_deserialize)
-    set_confchangev2_context_deserializer(pickle_deserialize)
-    set_entry_context_deserializer(pickle_deserialize)
-    set_entry_data_deserializer(pickle_deserialize)
-    set_message_context_deserializer(pickle_deserialize)
-    set_snapshot_data_deserializer(pickle_deserialize)
-    set_fsm_deserializer(pickle_deserialize)
-    set_log_entry_deserializer(pickle_deserialize)
-
+    set_custom_formatters(
+        entry_data=pickle_deserialize,
+        entry_context=pickle_deserialize,
+        confchange_context=pickle_deserialize,
+        confchangev2_context=pickle_deserialize,
+        message_context=pickle_deserialize,
+        snapshot_data=pickle_deserialize,
+        log_entry=pickle_deserialize,
+        fsm=pickle_deserialize,
+    )
 
 class HashStore:
     """

binding/python/examples/deserializer.py

@@ -1,12 +1,5 @@
 import pickle
-from raftify import (
-    set_confchange_context_deserializer,
-    set_confchangev2_context_deserializer,
-    set_entry_context_deserializer,
-    set_entry_data_deserializer,
-    set_message_context_deserializer,
-    set_snapshot_data_deserializer,
-)
+from raftify import set_custom_formatters
 
 
 def pickle_deserialize(data: bytes) -> str | None:
@@ -26,9 +19,13 @@ def register_custom_deserializer() -> None:
     Initialize the custom deserializers.
     """
 
-    set_confchange_context_deserializer(pickle_deserialize)
-    set_confchangev2_context_deserializer(pickle_deserialize)
-    set_entry_context_deserializer(pickle_deserialize)
-    set_entry_data_deserializer(pickle_deserialize)
-    set_message_context_deserializer(pickle_deserialize)
-    set_snapshot_data_deserializer(pickle_deserialize)
+    set_custom_formatters(
+        entry_data=pickle_deserialize,
+        entry_context=pickle_deserialize,
+        confchange_context=pickle_deserialize,
+        confchangev2_context=pickle_deserialize,
+        message_context=pickle_deserialize,
+        snapshot_data=pickle_deserialize,
+        log_entry=pickle_deserialize,
+        fsm=pickle_deserialize,
+    )

binding/python/raftify.pyi

@@ -305,36 +305,19 @@ class RaftServiceClient:
     async def debug_node(self) -> str:
         """ """
 
-def set_snapshot_data_deserializer(cb: Callable[[bytes], str | bytes | None]) -> None:
-    """ """
-
-def set_message_context_deserializer(cb: Callable[[bytes], str | bytes | None]) -> None:
-    """ """
-
-def set_confchange_context_deserializer(
-    cb: Callable[[bytes], str | bytes | None]
+def set_custom_formatters(
+    *,
+    entry_context: Optional[Callable[[bytes], str | bytes | None]] = None,
+    entry_data: Optional[Callable[[bytes], str | bytes | None]] = None,
+    confchangev2_context: Optional[Callable[[bytes], str | bytes | None]] = None,
+    confchange_context: Optional[Callable[[bytes], str | bytes | None]] = None,
+    message_context: Optional[Callable[[bytes], str | bytes | None]] = None,
+    snapshot_data: Optional[Callable[[bytes], str | bytes | None]] = None,
+    log_entry: Optional[Callable[[bytes], str | bytes | None]] = None,
+    fsm: Optional[Callable[[bytes], str | bytes | None]] = None,
 ) -> None:
     """ """
 
-def set_confchangev2_context_deserializer(
-    cb: Callable[[bytes], str | bytes | None]
-) -> None:
-    """ """
-
-def set_entry_data_deserializer(cb: Callable[[bytes], str | bytes | None]) -> None:
-    """ """
-
-def set_entry_context_deserializer(cb: Callable[[bytes], str | bytes | None]) -> None:
-    """ """
-
-def set_fsm_deserializer(cb: Callable[[bytes], str | bytes | None]) -> None:
-    """ """
-    ...
-
-def set_log_entry_deserializer(cb: Callable[[bytes], str | bytes | None]) -> None:
-    """ """
-    ...
-
 class ConfChangeTransition:
     """ """
 

binding/python/src/bindings/formatter.rs

@@ -21,34 +21,45 @@ static MESSAGE_CONTEXT_DESERIALIZER_CB: Lazy<Mutex<Option<PyObject>>> =
 static SNAPSHOT_DATA_DESERIALIZER_CB: Lazy<Mutex<Option<PyObject>>> =
     Lazy::new(|| Mutex::new(None));
 
-#[pyfunction]
-pub fn set_entry_context_deserializer(cb: PyObject) {
-    *ENTRY_CONTEXT_DESERIALIZE_CB.lock().unwrap() = Some(cb);
-}
-
-#[pyfunction]
-pub fn set_entry_data_deserializer(cb: PyObject) {
-    *ENTRY_DATA_DESERIALIZE_CB.lock().unwrap() = Some(cb);
-}
-
-#[pyfunction]
-pub fn set_confchangev2_context_deserializer(cb: PyObject) {
-    *CONFCHANGEV2_CONTEXT_DESERIALIZE_CB.lock().unwrap() = Some(cb);
-}
-
-#[pyfunction]
-pub fn set_confchange_context_deserializer(cb: PyObject) {
-    *CONFCHANGE_CONTEXT_DESERIALIZE_CB.lock().unwrap() = Some(cb);
-}
-
-#[pyfunction]
-pub fn set_message_context_deserializer(cb: PyObject) {
-    *MESSAGE_CONTEXT_DESERIALIZER_CB.lock().unwrap() = Some(cb);
-}
+pub static ENTRY_LOG_ENTRY_DESERIALIZE_CB: Lazy<Mutex<Option<PyObject>>> =
+    Lazy::new(|| Mutex::new(None));
+pub static ENTRY_FSM_DESERIALIZE_CB: Lazy<Mutex<Option<PyObject>>> = Lazy::new(|| Mutex::new(None));
 
 #[pyfunction]
-pub fn set_snapshot_data_deserializer(cb: PyObject) {
-    *SNAPSHOT_DATA_DESERIALIZER_CB.lock().unwrap() = Some(cb);
+pub fn set_custom_formatters(
+    entry_context: Option<PyObject>,
+    entry_data: Option<PyObject>,
+    confchangev2_context: Option<PyObject>,
+    confchange_context: Option<PyObject>,
+    message_context: Option<PyObject>,
+    snapshot_data: Option<PyObject>,
+    log_entry: Option<PyObject>,
+    fsm: Option<PyObject>,
+) {
+    if let Some(cb) = entry_context {
+        *ENTRY_CONTEXT_DESERIALIZE_CB.lock().unwrap() = Some(cb);
+    }
+    if let Some(cb) = entry_data {
+        *ENTRY_DATA_DESERIALIZE_CB.lock().unwrap() = Some(cb);
+    }
+    if let Some(cb) = confchangev2_context {
+        *CONFCHANGEV2_CONTEXT_DESERIALIZE_CB.lock().unwrap() = Some(cb);
+    }
+    if let Some(cb) = confchange_context {
+        *CONFCHANGE_CONTEXT_DESERIALIZE_CB.lock().unwrap() = Some(cb);
+    }
+    if let Some(cb) = message_context {
+        *MESSAGE_CONTEXT_DESERIALIZER_CB.lock().unwrap() = Some(cb);
+    }
+    if let Some(cb) = snapshot_data {
+        *SNAPSHOT_DATA_DESERIALIZER_CB.lock().unwrap() = Some(cb);
+    }
+    if let Some(cb) = log_entry {
+        *ENTRY_LOG_ENTRY_DESERIALIZE_CB.lock().unwrap() = Some(cb);
+    }
+    if let Some(cb) = fsm {
+        *ENTRY_FSM_DESERIALIZE_CB.lock().unwrap() = Some(cb);
+    }
 }
 
 impl CustomFormatter for PythonFormatter {

binding/python/src/bindings/state_machine.rs

@@ -1,29 +1,15 @@
 use async_trait::async_trait;
-use once_cell::sync::Lazy;
 use pyo3::{prelude::*, types::PyBytes};
 use pyo3_asyncio::TaskLocals;
 use raftify::{AbstractLogEntry, AbstractStateMachine, Error, Result};
-use std::{fmt, sync::Mutex};
+use std::fmt;
 
+use super::formatter::{ENTRY_FSM_DESERIALIZE_CB, ENTRY_LOG_ENTRY_DESERIALIZE_CB};
 use super::{
     errors::{ApplyError, RestoreError, SnapshotError},
     utils::get_python_repr,
 };
 
-pub static ENTRY_LOG_ENTRY_DESERIALIZE_CB: Lazy<Mutex<Option<PyObject>>> =
-    Lazy::new(|| Mutex::new(None));
-pub static ENTRY_FSM_DESERIALIZE_CB: Lazy<Mutex<Option<PyObject>>> = Lazy::new(|| Mutex::new(None));
-
-#[pyfunction]
-pub fn set_log_entry_deserializer(cb: PyObject) {
-    *ENTRY_LOG_ENTRY_DESERIALIZE_CB.lock().unwrap() = Some(cb);
-}
-
-#[pyfunction]
-pub fn set_fsm_deserializer(cb: PyObject) {
-    *ENTRY_FSM_DESERIALIZE_CB.lock().unwrap() = Some(cb);
-}
-
 #[derive(Clone)]
 #[pyclass]
 pub struct PyLogEntry {

binding/python/src/lib.rs

@@ -36,42 +36,7 @@ fn raftify(py: Python, m: &PyModule) -> PyResult<()> {
     m.add_function(wrap_pyfunction!(bindings::cli::cli_main, m)?)?;
 
     m.add_function(wrap_pyfunction!(
-        bindings::state_machine::set_log_entry_deserializer,
-        m
-    )?)?;
-
-    m.add_function(wrap_pyfunction!(
-        bindings::state_machine::set_fsm_deserializer,
-        m
-    )?)?;
-
-    m.add_function(wrap_pyfunction!(
-        bindings::formatter::set_confchange_context_deserializer,
-        m
-    )?)?;
-
-    m.add_function(wrap_pyfunction!(
-        bindings::formatter::set_confchangev2_context_deserializer,
-        m
-    )?)?;
-
-    m.add_function(wrap_pyfunction!(
-        bindings::formatter::set_entry_context_deserializer,
-        m
-    )?)?;
-
-    m.add_function(wrap_pyfunction!(
-        bindings::formatter::set_entry_data_deserializer,
-        m
-    )?)?;
-
-    m.add_function(wrap_pyfunction!(
-        bindings::formatter::set_message_context_deserializer,
-        m
-    )?)?;
-
-    m.add_function(wrap_pyfunction!(
-        bindings::formatter::set_snapshot_data_deserializer,
+        bindings::formatter::set_custom_formatters,
         m
     )?)?;