Query context manager (#461)

* remove useless `IntoRust` bound in handlers system

Only the `IntoPython` bound is important, the rest was there to
reinforce the typing, but it's not needed at all, and it prevented
to use objects not implementing `IntoRust` (which will be the case
when `Query` will use `option_wrapper`)

* Add context manager to queries, to properly use channels with queryables

* fix tests with new aligned queryable print

* add documentation for `Query.drop`

Author: wyfo

examples/z_queryable.py

@@ -17,32 +17,27 @@
 
 
 def main(conf: zenoh.Config, key: str, payload: str, complete: bool):
-    def queryable_callback(query):
-        print(
-            f">> [Queryable ] Received Query '{query.selector}'"
-            + (
-                f" with payload: {query.payload.to_string()}"
-                if query.payload is not None
-                else ""
-            )
-        )
-        query.reply(key, payload)
-
     # initiate logging
     zenoh.init_log_from_env_or("error")
 
     print("Opening session...")
     with zenoh.open(conf) as session:
         print(f"Declaring Queryable on '{key}'...")
-        session.declare_queryable(key, queryable_callback, complete=complete)
+        queryable = session.declare_queryable(key, complete=complete)
 
         print("Press CTRL-C to quit...")
         while True:
-            try:
-                time.sleep(1)
-            except Exception as err:
-                print(err, flush=True)
-                raise
+            with queryable.recv() as query:
+                if query.payload is not None:
+                    print(
+                        f">> [Queryable ] Received Query '{query.selector}'"
+                        f" with payload: '{query.payload.to_string()}'"
+                    )
+                else:
+                    print(f">> [Queryable ] Received Query '{query.selector}'")
+                query.reply(key, payload)
+                # it's possible to call `query.drop()` after handling it
+                # instead of using a context manager
 
 
 # --- Command line argument parsing --- --- --- --- --- ---

src/handlers.rs

@@ -207,6 +207,8 @@ impl Drop for PythonCallback {
     }
 }
 
+// the generic type is not useful per se, it just there to make typing
+// prettier, e.g. to have `get` returning a `PyResult<HandlerImpl<Reply>>`
 pub(crate) enum HandlerImpl<T> {
     Rust(Py<Handler>, PhantomData<T>),
     Python(PyObject),
@@ -239,10 +241,7 @@ impl<T> ToPyObject for HandlerImpl<T> {
     }
 }
 
-impl<T: IntoRust> HandlerImpl<T>
-where
-    T::Into: IntoPython,
-{
+impl<T> HandlerImpl<T> {
     pub(crate) fn try_recv(&self, py: Python) -> PyResult<PyObject> {
         match self {
             Self::Rust(handler, _) => handler.borrow(py).try_recv(py),
@@ -258,11 +257,11 @@ where
     }
 }
 
-struct RustHandler<H: IntoRust, T: IntoRust>
+struct RustHandler<H: IntoRust, T: IntoPython>
 where
-    H::Into: IntoHandler<T::Into>,
+    H::Into: IntoHandler<T>,
 {
-    handler: <H::Into as IntoHandler<T::Into>>::Handler,
+    handler: <H::Into as IntoHandler<T>>::Handler,
     _phantom: PhantomData<T>,
 }
 
@@ -300,10 +299,7 @@ impl<E: fmt::Display> fmt::Display for DeadlineError<E> {
     }
 }
 
-impl<T: IntoRust> Receiver for RustHandler<DefaultHandler, T>
-where
-    T::Into: IntoPython,
-{
+impl<T: IntoPython> Receiver for RustHandler<DefaultHandler, T> {
     fn type_name(&self) -> &'static str {
         short_type_name::<T>()
     }
@@ -325,10 +321,7 @@ where
     }
 }
 
-impl<T: IntoRust> Receiver for RustHandler<FifoChannel, T>
-where
-    T::Into: IntoPython,
-{
+impl<T: IntoPython> Receiver for RustHandler<FifoChannel, T> {
     fn type_name(&self) -> &'static str {
         short_type_name::<T>()
     }
@@ -350,10 +343,7 @@ where
     }
 }
 
-impl<T: IntoRust> Receiver for RustHandler<RingChannel, T>
-where
-    T::Into: IntoPython,
-{
+impl<T: IntoPython> Receiver for RustHandler<RingChannel, T> {
     fn type_name(&self) -> &'static str {
         short_type_name::<T>()
     }
@@ -375,14 +365,13 @@ where
     }
 }
 
-fn rust_handler<H: IntoRust, T: IntoRust>(
+fn rust_handler<H: IntoRust, T: IntoPython>(
     py: Python,
     into_handler: H,
-) -> (RustCallback<T::Into>, HandlerImpl<T>)
+) -> (RustCallback<T>, HandlerImpl<T::Into>)
 where
-    H::Into: IntoHandler<T::Into>,
-    <H::Into as IntoHandler<T::Into>>::Handler: Send + Sync,
-    T::Into: IntoPython,
+    H::Into: IntoHandler<T>,
+    <H::Into as IntoHandler<T>>::Handler: Send + Sync,
     RustHandler<H, T>: Receiver,
 {
     let (callback, handler) = into_handler.into_rust().into_handler();
@@ -418,13 +407,10 @@ fn python_callback<T: IntoPython>(callback: &Bound<PyAny>) -> PyResult<RustCallb
     })
 }
 
-pub(crate) fn into_handler<T: IntoRust>(
+pub(crate) fn into_handler<T: IntoPython>(
     py: Python,
     obj: Option<&Bound<PyAny>>,
-) -> PyResult<(impl IntoHandler<T::Into, Handler = HandlerImpl<T>>, bool)>
-where
-    T::Into: IntoPython,
-{
+) -> PyResult<(impl IntoHandler<T, Handler = HandlerImpl<T::Into>>, bool)> {
     let mut background = false;
     let Some(obj) = obj else {
         return Ok((rust_handler(py, DefaultHandler), background));

src/query.rs

@@ -76,38 +76,53 @@ impl QueryConsolidation {
     }
 }
 
-wrapper!(zenoh::query::Query: Clone);
+option_wrapper!(zenoh::query::Query, "Dropped query");
 
 #[pymethods]
 impl Query {
+    fn __enter__<'a, 'py>(this: &'a Bound<'py, Self>) -> &'a Bound<'py, Self> {
+        this
+    }
+
+    #[pyo3(signature = (*_args, **_kwargs))]
+    fn __exit__(
+        &mut self,
+        py: Python,
+        _args: &Bound<PyTuple>,
+        _kwargs: Option<&Bound<PyDict>>,
+    ) -> PyResult<PyObject> {
+        self.drop();
+        Ok(py.None())
+    }
+
     #[getter]
-    fn selector(&self) -> Selector {
-        self.0.selector().into_owned().into()
+    fn selector(&self) -> PyResult<Selector> {
+        Ok(self.get_ref()?.selector().into_owned().into())
     }
 
     #[getter]
-    fn key_expr(&self) -> KeyExpr {
-        self.0.key_expr().clone().into_owned().into()
+    fn key_expr(&self) -> PyResult<KeyExpr> {
+        Ok(self.get_ref()?.key_expr().clone().into_owned().into())
     }
 
     #[getter]
-    fn parameters(&self) -> Parameters {
-        self.0.parameters().clone().into_owned().into()
+    fn parameters(&self) -> PyResult<Parameters> {
+        Ok(self.get_ref()?.parameters().clone().into_owned().into())
     }
 
     #[getter]
-    fn payload(&self) -> Option<ZBytes> {
-        self.0.payload().cloned().map_into()
+    fn payload(&self) -> PyResult<Option<ZBytes>> {
+        Ok(self.get_ref()?.payload().cloned().map_into())
     }
 
     #[getter]
-    fn encoding(&self) -> Option<Encoding> {
-        self.0.encoding().cloned().map_into()
+    fn encoding(&self) -> PyResult<Option<Encoding>> {
+        Ok(self.get_ref()?.encoding().cloned().map_into())
     }
 
     #[getter]
-    fn attachment(&self) -> Option<ZBytes> {
-        self.0.attachment().cloned().map_into()
+    fn attachment(&self) -> PyResult<Option<ZBytes>> {
+        Ok(self.get_ref()?.attachment().cloned().map_into())
     }
 
     // TODO timestamp
@@ -125,7 +140,7 @@ impl Query {
         #[pyo3(from_py_with = "ZBytes::from_py_opt")] attachment: Option<ZBytes>,
     ) -> PyResult<()> {
         let build = build!(
-            self.0.reply(key_expr, payload),
+            self.get_ref()?.reply(key_expr, payload),
             encoding,
             congestion_control,
             priority,
@@ -134,14 +149,15 @@ impl Query {
         );
         wait(py, build)
     }
+
     #[pyo3(signature = (payload, *, encoding = None))]
     fn reply_err(
         &self,
         py: Python,
         #[pyo3(from_py_with = "ZBytes::from_py")] payload: ZBytes,
         #[pyo3(from_py_with = "Encoding::from_py_opt")] encoding: Option<Encoding>,
     ) -> PyResult<()> {
-        let build = build!(self.0.reply_err(payload), encoding);
+        let build = build!(self.get_ref()?.reply_err(payload), encoding);
         wait(py, build)
     }
 
@@ -156,7 +172,7 @@ impl Query {
         #[pyo3(from_py_with = "ZBytes::from_py_opt")] attachment: Option<ZBytes>,
     ) -> PyResult<()> {
         let build = build!(
-            self.0.reply_del(key_expr),
+            self.get_ref()?.reply_del(key_expr),
             congestion_control,
             priority,
             express,
@@ -165,12 +181,16 @@ impl Query {
         wait(py, build)
     }
 
-    fn __repr__(&self) -> String {
-        format!("{:?}", self.0)
+    fn drop(&mut self) {
+        Python::with_gil(|gil| gil.allow_threads(|| drop(self.0.take())));
     }
 
-    fn __str__(&self) -> String {
-        format!("{}", self.0)
+    fn __repr__(&self) -> PyResult<String> {
+        Ok(format!("{:?}", self.get_ref()?))
+    }
+
+    fn __str__(&self) -> PyResult<String> {
+        Ok(format!("{}", self.get_ref()?))
     }
 }
 

src/utils.rs

@@ -37,14 +37,14 @@ impl<T, E: IntoPyErr> IntoPyResult<T> for Result<T, E> {
     }
 }
 
-pub(crate) trait IntoRust: Send + Sync + 'static {
+pub(crate) trait IntoRust: 'static {
     type Into;
     fn into_rust(self) -> Self::Into;
 }
 
 into_rust!(bool, Duration);
 
-pub(crate) trait IntoPython: Sized + Send + 'static {
+pub(crate) trait IntoPython: Sized + Send + Sync + 'static {
     type Into: IntoPy<PyObject>;
     fn into_python(self) -> Self::Into;
     fn into_pyobject(self, py: Python) -> PyObject {

tests/examples_check.py

@@ -170,12 +170,12 @@ def test_z_querier_z_queryable():
         z_querier.errors.append("z_querier didn't get a response from z_queryable")
     queryableout = "".join(z_queryable.stdout)
     if not (
-        "Received Query 'demo/example/zenoh-python-queryable' with payload: [   0] value"
+        "Received Query 'demo/example/zenoh-python-queryable' with payload: '[   0] value'"
         in queryableout
     ):
         z_queryable.errors.append("z_queryable didn't catch query [0]")
     elif not (
-        "Received Query 'demo/example/zenoh-python-queryable' with payload: [   2] value"
+        "Received Query 'demo/example/zenoh-python-queryable' with payload: '[   2] value'"
         in queryableout
     ):
         z_queryable.errors.append("z_queryable didn't catch query [2]")

zenoh/__init__.pyi

@@ -509,6 +509,8 @@ class Publisher:
 class Query:
     """Structs received by a Queryable."""
 
+    def __enter__(self) -> Self: ...
+    def __exit__(self, *_args, **_kwargs): ...
     @property
     def selector(self) -> Selector: ...
     @property
@@ -552,6 +554,19 @@ class Query:
         By default, queries only accept replies whose key expression intersects with the query's. Unless the query has enabled disjoint replies (you can check this through Query::accepts_replies), replying on a disjoint key expression will result in an error when resolving the reply.
         """
 
+    def drop(self):
+        """Drop the instance of a query.
+        The query will only be finalized when all query instances (one per queryable
+        matched) are dropped. Finalization is required to not have query hanging
+        on the querier side.
+
+        This method should be called after handling the query, as Python finalizers
+        are not reliable, especially when it comes to loop variables. It is also
+        possible, and advised, to use query context manager, which calls `drop` on
+        exit. Once a query is dropped, it's no more possible to use it, and its
+        methods will raise an exception.
+        """
+
     def __str__(self) -> str: ...
 
 @final