json state proxy

Cargo.lock

@@ -795,6 +795,7 @@ dependencies = [
  "pyo3-log",
  "pyo3-stub-gen",
  "rand 0.10.1",
+ "serde_json",
  "tokio",
  "zenoh",
  "zerompk",

rust/exo_pyo3_bindings/Cargo.toml

@@ -60,3 +60,4 @@ env_logger.workspace = true
 zenoh.workspace = true
 zerompk.workspace = true
 rand = "0.10.1"
+serde_json = "1.0.149"

rust/exo_pyo3_bindings/exo_pyo3_bindings.pyi

@@ -1,72 +1,11 @@
 # This file is automatically generated by pyo3_stub_gen
 # ruff: noqa: E501, F401
 
-import builtins
 import typing
 
 @typing.final
-class Keypair:
-    r"""
-    Identity keypair of a node.
-    """
+class StateProxy:
     @staticmethod
-    def generate() -> Keypair:
-        r"""
-        Generate a new Ed25519 keypair.
-        """
-    @staticmethod
-    def from_bytes(bytes: bytes) -> Keypair:
-        r"""
-        Construct an Ed25519 keypair from secret key bytes
-        """
-    def to_bytes(self) -> bytes:
-        r"""
-        Get the secret key bytes underlying the keypair
-        """
-    def to_node_id(self) -> builtins.str:
-        r"""
-        Convert the `Keypair` into the corresponding `PeerId` string, which we use as our `NodeId`.
-        """
-
-@typing.final
-class NetworkingHandle:
-    def __new__(cls, identity: Keypair, bootstrap_peers: typing.Sequence[builtins.str], listen_port: builtins.int) -> NetworkingHandle: ...
-    async def gossipsub_subscribe(self, topic: builtins.str) -> builtins.bool:
-        r"""
-        Subscribe to a `GossipSub` topic.
-        
-        Returns `True` if the subscription worked. Returns `False` if we were already subscribed.
-        """
-    async def gossipsub_unsubscribe(self, topic: builtins.str) -> builtins.bool:
-        r"""
-        Unsubscribes from a `GossipSub` topic.
-        
-        Returns `True` if we were subscribed to this topic. Returns `False` if we were not subscribed.
-        """
-    async def gossipsub_publish(self, topic: builtins.str, data: bytes) -> None:
-        r"""
-        Publishes a message with multiple topics to the `GossipSub` network.
-        
-        If no peers are found that subscribe to this topic, throws `NoPeersSubscribedToTopicError` exception.
-        """
-    async def recv(self) -> PyFromSwarm: ...
-
-class PyFromSwarm:
-    @typing.final
-    class Connection(PyFromSwarm):
-        __match_args__ = ("connected",)
-        @property
-        def connected(self) -> builtins.bool: ...
-        def __new__(cls, connected: builtins.bool) -> PyFromSwarm.Connection: ...
-    
-    @typing.final
-    class Message(PyFromSwarm):
-        __match_args__ = ("topic", "data",)
-        @property
-        def topic(self) -> builtins.str: ...
-        @property
-        def data(self) -> bytes: ...
-        def __new__(cls, topic: builtins.str, data: bytes) -> PyFromSwarm.Message: ...
-    
-    ...
+    async def init() -> StateProxy: ...
+    async def snapshot(self) -> str: ...
 

rust/exo_pyo3_bindings/src/lib.rs

@@ -5,21 +5,15 @@
 //!
 
 mod allow_threading;
-mod ident;
-mod networking;
+// mod ident;
+// mod networking;
+mod state;
 
-use crate::ident::PyKeypair;
-use crate::networking::networking_submodule;
+use crate::state::snapshot_module;
 use pyo3::prelude::PyModule;
-use pyo3::types::PyModuleMethods;
-use pyo3::{Bound, PyResult, pyclass, pymodule};
+use pyo3::{Bound, PyResult, pymodule};
 use pyo3_stub_gen::define_stub_info_gatherer;
 
-/// Namespace for all the constants used by this crate.
-pub(crate) mod r#const {
-    pub const MPSC_CHANNEL_SIZE: usize = 1024;
-}
-
 /// Namespace for crate-wide extension traits/methods
 pub(crate) mod ext {
     use crate::allow_threading::AllowThreads;
@@ -162,8 +156,9 @@ fn main_module(m: &Bound<'_, PyModule>) -> PyResult<()> {
     // TODO: for now this is all NOT a submodule, but figure out how to make the submodule system
     //       work with maturin, where the types generate correctly, in the right folder, without
     //       too many importing issues...
-    m.add_class::<PyKeypair>()?;
-    networking_submodule(m)?;
+    // m.add_class::<PyKeypair>()?;
+    // networking_submodule(m)?;
+    snapshot_module(m)?;
 
     // top-level constructs
     // TODO: ...

rust/exo_pyo3_bindings/src/state.rs

@@ -0,0 +1,110 @@
+use networking::Session;
+use pyo3::{
+    exceptions::{PyRuntimeError, PyValueError},
+    prelude::*,
+};
+use pyo3_stub_gen::derive::{gen_methods_from_python, gen_stub_pyclass, gen_stub_pymethods};
+use serde_json::{Map, Value};
+use zenoh::{Result, Session as ZSession, query::QueryTarget, sample::SampleFields};
+
+pyo3_stub_gen::inventory::submit! {
+    gen_methods_from_python! {
+        r#"
+            class StateProxy:
+                @staticmethod
+                async def init() -> StateProxy: ...
+                async def snapshot(self) -> str: ...
+        "#
+    }
+}
+
+pub fn snapshot_module(m: &Bound<'_, PyModule>) -> PyResult<()> {
+    m.add_class::<StateProxy>()?;
+
+    Ok(())
+}
+
+#[gen_stub_pyclass]
+#[pyclass]
+pub struct StateProxy {
+    session: Session,
+}
+#[gen_stub_pymethods]
+#[pymethods]
+impl StateProxy {
+    #[staticmethod]
+    #[gen_stub(skip)]
+    pub fn init<'py>(py: Python<'py>) -> PyResult<Bound<'py, PyAny>> {
+        pyo3_async_runtimes::tokio::future_into_py(py, async move {
+            Ok(Self {
+                session: networking::open(
+                    networking::cfg(rand::random(), 0).expect("default cfg is valid"),
+                )
+                .await
+                .map_err(|e| PyRuntimeError::new_err(e.to_string()))?,
+            })
+        })
+    }
+
+    #[gen_stub(skip)]
+    pub fn snapshot<'py>(&self, py: Python<'py>) -> PyResult<Bound<'py, PyAny>> {
+        pyo3_async_runtimes::tokio::future_into_py(py, {
+            let session = self.session.clone();
+            async move {
+                Self::_snapshot(session)
+                    .await
+                    .map_err(|e| PyValueError::new_err(e.to_string()))
+                    .map(|v| v.to_string())
+            }
+        })
+    }
+}
+
+impl StateProxy {
+    async fn _snapshot(session: ZSession) -> Result<Value> {
+        let q = session.get("storage/mem1/**").await?;
+
+        let mut v = Value::Object(Map::default());
+
+        while let Ok(sample) = q.recv_async().await {
+            let mut cur_v = &mut v;
+            let Ok(sample) = sample.into_result() else {
+                continue;
+            };
+            // skip storage/mem1
+            let SampleFields {
+                payload, key_expr, ..
+            } = sample.into();
+            let mut iter = key_expr.split('/').skip(2).peekable();
+            loop {
+                let Some(p) = iter.next() else {
+                    break;
+                };
+                if iter.peek().is_none() {
+                    // terminal; write value into json
+                    let existing = cur_v
+                        .as_object_mut()
+                        .expect("path terminated unexpectedly - value stored at some/path and some/path/two")
+                        .insert(p.to_owned(), Value::String(payload.try_to_string()?.to_string()));
+
+                    if let Some(value) = existing {
+                        assert!(value.is_string())
+                        // could log, but string overwrites are fine
+                    }
+                } else {
+                    // non-terminal; ensure key exists in v, then replace cur with that object
+                    cur_v = cur_v
+                        .as_object_mut()
+                        .expect("path terminated unexpectedly - value stored at some/path and some/path/two")
+                        .entry(p)
+                        .or_insert(Value::Object(Map::default()));
+                    assert!(
+                        cur_v.is_object(),
+                        "path terminated unexpectedly - value stored at some/path and some/path/two"
+                    )
+                }
+            }
+        }
+        Ok(v)
+    }
+}

rust/networking/src/lib.rs

@@ -2,9 +2,11 @@ use std::{
     env,
     ops::{Deref, DerefMut},
     panic,
+    sync::Arc,
 };
 
 use netwatcher::WatchHandle;
+use parking_lot::Mutex;
 use tokio::{sync::mpsc, task::JoinHandle};
 use zenoh::{Result, Session as ZSession, config::WhatAmI, internal::runtime::Runtime};
 use zenoh_plugin_storage_manager::StoragesPlugin;
@@ -51,7 +53,7 @@ pub async fn open(cfg: zenoh::Config) -> Result<Session> {
         .await?;
     let session = zenoh::session::init(runtime.clone().into()).await?;
     runtime.start().await?;
-    let _watch_all_handle = watch_all(runtime).await?;
+    let _watch_all_handle = Arc::new(Mutex::new(watch_all(runtime).await?));
     Ok(Session {
         session,
         _watch_all_handle,
@@ -91,7 +93,7 @@ async fn watch_all(runtime: Runtime) -> Result<WatchAllHandle> {
 
 pub struct Session {
     pub session: ZSession,
-    _watch_all_handle: WatchAllHandle,
+    _watch_all_handle: Arc<Mutex<WatchAllHandle>>,
 }
 impl Deref for Session {
     type Target = ZSession;

tmp/state_proxy.py

@@ -0,0 +1,13 @@
+from exo_pyo3_bindings import StateProxy
+import anyio
+
+async def main():
+    sp = await StateProxy.init()
+    while True:
+        data = await sp.snapshot()
+        if data != "{}":
+            print(data)
+        await anyio.sleep(1)
+    
+if __name__ == "__main__":
+    anyio.run(main)