Added push and pull consumers.

Author: s3rius

python/natsrpy/_natsrpy_rs/js/consumers.pyi

@@ -147,7 +147,16 @@ class PushConsumerConfig:
         pause_until: int | None = None,
     ) -> None: ...
 
-class PushConsumer: ...
+class MessagesIterator:
+    def __aiter__(self) -> MessagesIterator: ...
+    async def __anext__(self) -> JetStreamMessage: ...
+    async def next(
+        self,
+        timeout: float | timedelta | None = None,
+    ) -> JetStreamMessage: ...
+
+class PushConsumer:
+    async def messages(self) -> MessagesIterator: ...
 
 class PullConsumer:
     async def fetch(

src/exceptions/rust_err.rs

@@ -1,4 +1,4 @@
-use pyo3::exceptions::{PyTimeoutError, PyTypeError};
+use pyo3::exceptions::{PyStopAsyncIteration, PyTimeoutError, PyTypeError};
 
 use crate::exceptions::py_err::{NatsrpyPublishError, NatsrpySessionError};
 
@@ -12,6 +12,8 @@ pub enum NatsrpyError {
     InvalidArgument(String),
     #[error("Session is not initialized. Call startup() first.")]
     NotInitialized,
+    #[error("The end of stream")]
+    AsyncStopIteration,
     #[error("Connection is closed or lost.")]
     Disconnected,
     #[error(transparent)]
@@ -66,12 +68,17 @@ pub enum NatsrpyError {
     PullConsumerError(#[from] async_nats::jetstream::stream::ConsumerError),
     #[error(transparent)]
     PullConsumerBatchError(#[from] async_nats::jetstream::consumer::pull::BatchError),
+    #[error(transparent)]
+    PushConsumerMessageError(#[from] async_nats::jetstream::consumer::push::MessagesError),
+    #[error(transparent)]
+    ConsumerStreamError(#[from] async_nats::jetstream::consumer::StreamError),
 }
 
 impl From<NatsrpyError> for pyo3::PyErr {
     fn from(value: NatsrpyError) -> Self {
         match value {
             NatsrpyError::PublishError(_) => NatsrpyPublishError::new_err(value.to_string()),
+            NatsrpyError::AsyncStopIteration => PyStopAsyncIteration::new_err("End of the stream."),
             NatsrpyError::Timeout(_) => PyTimeoutError::new_err(value.to_string()),
             NatsrpyError::PyError(py_err) => py_err,
             NatsrpyError::InvalidArgument(descr) => PyTypeError::new_err(descr),

src/js/consumers/mod.rs

@@ -9,5 +9,7 @@ pub mod pymod {
     #[pymodule_export]
     pub use super::pull::{config::PullConsumerConfig, consumer::PullConsumer};
     #[pymodule_export]
-    pub use super::push::{config::PushConsumerConfig, consumer::PushConsumer};
+    pub use super::push::{
+        config::PushConsumerConfig, consumer::MessagesIterator, consumer::PushConsumer,
+    };
 }

src/js/consumers/pull/consumer.rs

@@ -89,7 +89,7 @@ impl PullConsumer {
             let mut ret_messages = Vec::new();
             while let Some(msg) = messages.next().await {
                 let raw_msg = msg?;
-                ret_messages.push(crate::js::message::JetStreamMessage::from(raw_msg));
+                ret_messages.push(crate::js::message::JetStreamMessage::try_from(raw_msg)?);
             }
             Ok(ret_messages)
         })

src/js/consumers/push/consumer.rs

@@ -1,13 +1,20 @@
 use std::sync::Arc;
 
+use futures_util::StreamExt;
+use pyo3::{Bound, PyAny, PyRef, Python};
 use tokio::sync::RwLock;
 
+use crate::{
+    exceptions::rust_err::{NatsrpyError, NatsrpyResult},
+    js::pymod::JetStreamMessage,
+    utils::{futures::natsrpy_future_with_timeout, natsrpy_future, py_types::TimeValue},
+};
+
 type NatsPushConsumer =
     async_nats::jetstream::consumer::Consumer<async_nats::jetstream::consumer::push::Config>;
 
 #[pyo3::pyclass(from_py_object)]
 #[derive(Debug, Clone)]
-#[allow(dead_code)] // TODO! remove later.
 pub struct PushConsumer {
     consumer: Arc<RwLock<NatsPushConsumer>>,
 }
@@ -21,5 +28,67 @@ impl PushConsumer {
     }
 }
 
+#[pyo3::pyclass]
+pub struct MessagesIterator {
+    messages: Option<Arc<RwLock<async_nats::jetstream::consumer::push::Messages>>>,
+}
+
+impl From<async_nats::jetstream::consumer::push::Messages> for MessagesIterator {
+    fn from(value: async_nats::jetstream::consumer::push::Messages) -> Self {
+        Self {
+            messages: Some(Arc::new(RwLock::new(value))),
+        }
+    }
+}
+
+#[pyo3::pymethods]
+impl PushConsumer {
+    pub fn messages<'py>(&self, py: Python<'py>) -> NatsrpyResult<Bound<'py, PyAny>> {
+        let consumer_guard = self.consumer.clone();
+        natsrpy_future(py, async move {
+            Ok(MessagesIterator::from(
+                consumer_guard.read().await.messages().await?,
+            ))
+        })
+    }
+}
+
 #[pyo3::pymethods]
-impl PushConsumer {}
+impl MessagesIterator {
+    #[must_use]
+    pub const fn __aiter__(slf: PyRef<Self>) -> PyRef<Self> {
+        slf
+    }
+
+    pub fn next<'py>(
+        &self,
+        py: Python<'py>,
+        timeout: Option<TimeValue>,
+    ) -> NatsrpyResult<Bound<'py, PyAny>> {
+        let Some(messages_guard) = self.messages.clone() else {
+            unreachable!("Message is always Some in runtime.")
+        };
+        #[allow(clippy::significant_drop_tightening)]
+        natsrpy_future_with_timeout(py, timeout, async move {
+            let mut messages = messages_guard.write().await;
+            let Some(message) = messages.next().await else {
+                return Err(NatsrpyError::AsyncStopIteration);
+            };
+            let message = message?;
+
+            JetStreamMessage::try_from(message)
+        })
+    }
+
+    pub fn __anext__<'py>(&self, py: Python<'py>) -> NatsrpyResult<Bound<'py, PyAny>> {
+        self.next(py, None)
+    }
+}
+
+impl Drop for MessagesIterator {
+    fn drop(&mut self) {
+        pyo3_async_runtimes::tokio::get_runtime().block_on(async move {
+            self.messages = None;
+        });
+    }
+}

src/js/message.rs

@@ -1,27 +1,30 @@
-use pyo3::{Bound, Py, PyAny, Python, types::PyDict};
+use pyo3::{
+    Bound, Py, PyAny, Python,
+    types::{PyBytes, PyDict},
+};
 use std::sync::Arc;
 use tokio::sync::RwLock;
 
 use crate::{
-    exceptions::rust_err::NatsrpyResult,
-    utils::{headers::NatsrpyHeadermapExt, natsrpy_future, py_types::TimeValue},
+    exceptions::rust_err::{NatsrpyError, NatsrpyResult},
+    utils::{natsrpy_future, py_types::TimeValue},
 };
 
 #[pyo3::pyclass]
 pub struct JetStreamMessage {
-    message: async_nats::Message,
-    headers: Option<Py<PyDict>>,
+    message: crate::message::Message,
     acker: Arc<RwLock<async_nats::jetstream::message::Acker>>,
 }
 
-impl From<async_nats::jetstream::Message> for JetStreamMessage {
-    fn from(value: async_nats::jetstream::Message) -> Self {
+impl TryFrom<async_nats::jetstream::Message> for JetStreamMessage {
+    type Error = NatsrpyError;
+
+    fn try_from(value: async_nats::jetstream::Message) -> Result<Self, Self::Error> {
         let (message, acker) = value.split();
-        Self {
-            message,
-            headers: None,
+        Ok(Self {
+            message: message.try_into()?,
             acker: Arc::new(RwLock::new(acker)),
-        }
+        })
     }
 }
 
@@ -47,27 +50,23 @@ impl JetStreamMessage {
 #[pyo3::pymethods]
 impl JetStreamMessage {
     #[getter]
-    pub fn subject(&self) -> &str {
+    #[must_use]
+    pub const fn subject(&self) -> &str {
         self.message.subject.as_str()
     }
     #[getter]
-    pub fn reply(&self) -> Option<&str> {
-        self.message.reply.as_ref().map(async_nats::Subject::as_str)
+    #[must_use]
+    pub const fn reply(&self) -> &Option<String> {
+        &self.message.reply
     }
     #[getter]
-    pub fn payload(&self) -> &[u8] {
+    #[must_use]
+    pub const fn payload(&self) -> &Py<PyBytes> {
         &self.message.payload
     }
     #[getter]
-    pub fn headers(&mut self, py: Python<'_>) -> NatsrpyResult<Py<PyDict>> {
-        if let Some(headers) = &self.headers {
-            Ok(headers.clone_ref(py))
-        } else {
-            let headermap = self.message.headers.clone().unwrap_or_default();
-            let headers = headermap.to_pydict(py)?.unbind();
-            self.headers = Some(headers.clone_ref(py));
-            Ok(headers)
-        }
+    pub const fn headers(&mut self) -> &Py<PyDict> {
+        &self.message.headers
     }
 
     #[pyo3(signature=(double=false))]
@@ -107,4 +106,9 @@ impl JetStreamMessage {
     pub fn term<'py>(&self, py: Python<'py>, double: bool) -> NatsrpyResult<Bound<'py, PyAny>> {
         self.inner_ack(py, async_nats::jetstream::message::AckKind::Term, double)
     }
+
+    #[must_use]
+    pub fn __repr__(&self) -> String {
+        self.message.__repr__()
+    }
 }

src/subscription.rs

@@ -1,13 +1,12 @@
 use futures_util::StreamExt;
-use pyo3::exceptions::PyStopAsyncIteration;
-use std::{sync::Arc, time::Duration};
+use std::sync::Arc;
 
 use pyo3::{Bound, PyAny, PyRef, Python};
 use tokio::sync::Mutex;
 
 use crate::{
     exceptions::rust_err::{NatsrpyError, NatsrpyResult},
-    utils::futures::natsrpy_future_with_timeout,
+    utils::{futures::natsrpy_future_with_timeout, py_types::TimeValue},
 };
 
 #[pyo3::pyclass]
@@ -34,14 +33,14 @@ impl Subscription {
     pub fn next<'py>(
         &self,
         py: Python<'py>,
-        timeout: Option<Duration>,
+        timeout: Option<TimeValue>,
     ) -> NatsrpyResult<Bound<'py, PyAny>> {
         let Some(inner) = self.inner.clone() else {
             return Err(NatsrpyError::NotInitialized);
         };
         natsrpy_future_with_timeout(py, timeout, async move {
             let Some(message) = inner.lock().await.next().await else {
-                return Err(PyStopAsyncIteration::new_err("End of the stream.").into());
+                return Err(NatsrpyError::AsyncStopIteration);
             };
 
             crate::message::Message::try_from(message)