feat: implement media download functionality and enhance TryxClient methods

Author: krypton-byte

python/tryx/client.pyi

@@ -19,10 +19,19 @@ class Tryx:
     def run(self) -> Awaitable[None]: ...
     def run_blocking(self) -> None: ...
 
+DownloadableMedia = (
+    MessageProto.ImageMessage
+    | MessageProto.VideoMessage
+    | MessageProto.AudioMessage
+    | MessageProto.DocumentMessage
+    | MessageProto.StickerMessage
+)
+
 class TryxClient:
     async def send_message(self, chat: JID, message: MessageProto) -> str: ...
     async def upload(self, data: bytes, media_type: MediaType) -> UploadResponse: ...
     async def upload_file(self, path: str, media_type: MediaType) -> UploadResponse: ...
+    async def download_media(self, message: DownloadableMedia) -> bytes: ...
 
 class Nu[T]:
     X: T

src/client.rs

@@ -15,6 +15,7 @@ use whatsapp_rust::store::Backend;
 use whatsapp_rust_tokio_transport::TokioWebSocketTransportFactory;
 use whatsapp_rust_ureq_http_client::UreqHttpClient;
 use waproto::whatsapp::Message as WhatsappMessage;
+use waproto::whatsapp::message::{self as wa};
 use prost::Message;
 use tokio::signal;
 use tracing::{debug, error, info, warn};
@@ -55,13 +56,96 @@ pub struct TryxClient {
 
 #[pymethods]
 impl TryxClient {
+    fn is_connected(&self) -> bool {
+        self.client_rx.borrow().is_some()
+    }
+    fn download_media<'py>(&self, py: Python<'py>, message: Py<PyAny>) -> PyResult<Bound<'py, PyAny>> {
+        let client = self.client_rx.borrow().clone().ok_or_else(|| {
+            PyErr::new::<pyo3::exceptions::PyRuntimeError, _>("Bot is not running")
+        })?;
+        let message_type_name = message
+            .getattr(py, "DESCRIPTOR")
+            .and_then(|descriptor| descriptor.getattr(py, "name"))
+            .and_then(|name| name.extract::<String>(py))
+            .unwrap_or_default();
+        let serialized: Vec<u8> = message
+            .call_method0(py, "SerializeToString")?
+            .extract(py)?;
+
+        let locals = get_current_locals(py)?;
+        future_into_py_with_locals(py, locals, async move {
+            let download = match message_type_name.as_str() {
+                "ImageMessage" => {
+                    let media = wa::ImageMessage::decode(serialized.as_slice()).map_err(|e| {
+                        PyErr::new::<pyo3::exceptions::PyValueError, _>(
+                            format!("Failed to decode ImageMessage: {}", e),
+                        )
+                    })?;
+                    client.download(&media).await
+                }
+                "VideoMessage" => {
+                    let media = wa::VideoMessage::decode(serialized.as_slice()).map_err(|e| {
+                        PyErr::new::<pyo3::exceptions::PyValueError, _>(
+                            format!("Failed to decode VideoMessage: {}", e),
+                        )
+                    })?;
+                    client.download(&media).await
+                }
+                "DocumentMessage" => {
+                    let media = wa::DocumentMessage::decode(serialized.as_slice()).map_err(|e| {
+                        PyErr::new::<pyo3::exceptions::PyValueError, _>(
+                            format!("Failed to decode DocumentMessage: {}", e),
+                        )
+                    })?;
+                    client.download(&media).await
+                }
+                "AudioMessage" => {
+                    let media = wa::AudioMessage::decode(serialized.as_slice()).map_err(|e| {
+                        PyErr::new::<pyo3::exceptions::PyValueError, _>(
+                            format!("Failed to decode AudioMessage: {}", e),
+                        )
+                    })?;
+                    client.download(&media).await
+                }
+                "StickerMessage" => {
+                    let media = wa::StickerMessage::decode(serialized.as_slice()).map_err(|e| {
+                        PyErr::new::<pyo3::exceptions::PyValueError, _>(
+                            format!("Failed to decode StickerMessage: {}", e),
+                        )
+                    })?;
+                    client.download(&media).await
+                }
+                _ => {
+                    // Fallback path for unknown wrappers from Python side.
+                    if let Ok(media) = wa::ImageMessage::decode(serialized.as_slice()) {
+                        client.download(&media).await
+                    } else if let Ok(media) = wa::VideoMessage::decode(serialized.as_slice()) {
+                        client.download(&media).await
+                    } else if let Ok(media) = wa::DocumentMessage::decode(serialized.as_slice()) {
+                        client.download(&media).await
+                    } else if let Ok(media) = wa::AudioMessage::decode(serialized.as_slice()) {
+                        client.download(&media).await
+                    } else if let Ok(media) = wa::StickerMessage::decode(serialized.as_slice()) {
+                        client.download(&media).await
+                    } else {
+                        return Err(PyErr::new::<pyo3::exceptions::PyValueError, _>(
+                            "Failed to decode message as supported media message",
+                        ));
+                    }
+                }
+            };
+
+            download.map_err(|e| PyErr::new::<pyo3::exceptions::PyRuntimeError, _>(e.to_string()))
+        })
+    }
     fn upload_file<'py>(&self, py: Python<'py>, path: String, media_type: Py<MediaType>) -> PyResult<Bound<'py, PyAny>> {
         let client = self.client_rx.borrow().clone().ok_or_else(|| {
             PyErr::new::<pyo3::exceptions::PyRuntimeError, _>("Bot is not running")
         })?;
         let media_type_enum = media_type.bind(py).borrow_mut().to_wacore_enum();
         let locals = get_current_locals(py)?;
-        let data = std::fs::read(path.clone()).map_err(|e| PyErr::new::<pyo3::exceptions::PyRuntimeError, _>(e.to_string()))?;
+        let data = std::fs::read(&path)
+            .map_err(|e| PyErr::new::<pyo3::exceptions::PyRuntimeError, _>(e.to_string()))?;
         future_into_py_with_locals(py, locals, async move {
             let url = client
                 .upload(data, media_type_enum)

src/main.rs

@@ -11,27 +11,27 @@ async fn main() -> Result<(), Box<dyn std::error::Error>> {
     let backend = Arc::new(SqliteStore::new("whatsapp.db").await?);
 
     // Build the bot
-    let mut bot = Bot::builder()
+    let bot = Bot::builder()
         .with_backend(backend)
         .with_transport_factory(TokioWebSocketTransportFactory::new())
         .with_http_client(UreqHttpClient::new())
-        .on_event(|event, client| async move {
+        .on_event(|event, _client| async move {
             match event {
                 Event::PairingQrCode { code, .. } => {
                     println!("Scan this QR code with WhatsApp:\n{}", code);
                 }
                 Event::Message(msg, info) => {
                     println!("Message from {}: {:?}", info.source.sender, msg);
                 }
-                 Event::Connected(e)=> {
+                 Event::Connected(_e)=> {
                     println!("Connected to WhatsApp");
                 }
                 _ => {}
             }
         })
         .build();
         let mut bot2 = bot.await?;
-        let g = bot2.client();
+        let _g = bot2.client();
 
     // Start the bot
     bot2.run().await?.await?;

src/types.rs

@@ -1,9 +1,8 @@
 use std::sync::Arc;
 
-use pyo3::{PyAny, PyErr, PyResult, Python, exceptions::{PyException, PyRuntimeError}, ffi::PyObject, pyclass, pymethods, types::{PyAnyMethods, PyBytes, PyDateTime, PyType}};
-use waproto::whatsapp::VerifiedNameCertificate;
+use pyo3::{PyAny, PyErr, PyResult, Python, exceptions::PyRuntimeError, pyclass, pymethods, types::{PyAnyMethods, PyBytes, PyDateTime}};
 use whatsapp_rust::{Jid as WhatsAppJID};
-use wacore::types::message::{BotEditType, EditAttribute, MessageInfo as WhatsAppMessageInfo, MessageSource as WhatsAppMessageSource, MsgBotInfo as WhatsAppMsgBotInfo, MsgMetaInfo as WhatsappMsgMetaInfo};
+use wacore::types::message::{BotEditType, EditAttribute, MessageInfo as WhatsAppMessageInfo, MessageSource as WhatsAppMessageSource, MsgBotInfo as WhatsAppMsgBotInfo};
 use prost::Message;