feat: add upload functionality and MediaType enum for media handling

Author: krypton-byte

python/tryx/client.pyi

@@ -4,8 +4,9 @@ from typing import Any, Awaitable, Callable, Type
 
 # from tryx.events import Message
 from .waproto.whatsapp_pb2 import Message  as MessageProto
-from .types import JID
+from .types import JID, UploadResponse
 from .backend import SqliteBackend
+from .wacore import MediaType
 
 class Tryx:
     def __init__(self, backend: SqliteBackend) -> None: ...
@@ -17,6 +18,9 @@ class Tryx:
 
 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: ...
+
 
 class Nu[T]:
     X: T

python/tryx/types.pyi

@@ -60,3 +60,16 @@ class MessageInfo:
     @property
     def device_sent_meta(self) -> DeviceSentMeta | None: ...
 
+class UploadResponse:
+    @property
+    def url(self) -> str: ...
+    @property
+    def direct_path(self) -> str: ...
+    @property
+    def media_key(self) -> bytes: ...
+    @property
+    def file_enc_sha256(self) -> str: ...
+    @property
+    def file_sha256(self) -> str: ...
+    @property
+    def file_length(self) -> int: ...

python/tryx/wacore.py

@@ -0,0 +1,8 @@
+from ._tryx import wacore # type: ignore
+
+for name in dir(wacore): # type: ignore
+    obj = getattr(wacore, name) # type: ignore
+    if isinstance(obj, type):
+        globals()[name] = obj
+
+__all__ = [name for name in dir(wacore) if isinstance(getattr(wacore, name), type)] # type: ignore

python/tryx/wacore.pyi

@@ -0,0 +1,13 @@
+from enum import Enum
+
+
+class MediaType(Enum):
+    Image = 1
+    Video = 2
+    Audio = 3
+    Document = 4
+    History = 5
+    AppState = 6
+    Sticker = 7
+    StickerPack = 8
+    LinkThumbnail = 9

src/client.rs

@@ -21,10 +21,11 @@ use tracing::{debug, error, info, warn};
 use tracing_subscriber::EnvFilter;
 
 use crate::backend::{SqliteBackend, BackendBase};
-use crate::events::{Connected, Message as WAMessage, PairingQrCode};
+use crate::events::{Connected, LoggedOut, Message as WAMessage, PairingQrCode};
 use crate::exceptions::UnsupportedBackend;
-use crate::types::JID;
+use crate::types::{JID, UploadResponse};
 use crate::dispatcher::Dispatcher;
+use crate::wacore::MediaType;
 
 static LOG_INIT: Once = Once::new();
 
@@ -54,6 +55,54 @@ pub struct TryxClient {
 
 #[pymethods]
 impl TryxClient {
+    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()))?;
+        future_into_py_with_locals(py, locals, async move {
+            let url = client
+                .upload(data, media_type_enum)
+                .await
+                .map_err(|e| PyErr::new::<pyo3::exceptions::PyRuntimeError, _>(e.to_string()))?;
+            let result= UploadResponse {
+                url: url.url,
+                direct_path: url.direct_path,
+                media_key: url.media_key,
+                file_enc_sha256: url.file_enc_sha256,
+                file_sha256: url.file_sha256,
+                file_length: url.file_length,
+            };
+            Ok(result)
+        })
+    }
+    fn upload<'py>(&self, py: Python<'py>, data: &[u8], 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 data_vec = data.to_vec();
+        let mtype = media_type.bind(py).borrow_mut().to_wacore_enum();
+        let locals = get_current_locals(py)?;
+        future_into_py_with_locals::<_, UploadResponse>(py, locals, async move {
+            let url = client
+                .upload(data_vec, mtype)
+                .await
+                .map_err(|e| PyErr::new::<pyo3::exceptions::PyRuntimeError, _>(e.to_string()))?;
+            let result= UploadResponse {
+                url: url.url,
+                direct_path: url.direct_path,
+                media_key: url.media_key,
+                file_enc_sha256: url.file_enc_sha256,
+                file_sha256: url.file_sha256,
+                file_length: url.file_length,
+            };
+            Ok(result)
+        })
+        //     Ok(url)
+        // })
+    }
     fn send_message<'py>(&self, py: Python<'py>, to: Py<JID>, 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")
@@ -91,28 +140,47 @@ impl Tryx {
         tryx_client: Py<TryxClient>,
         client_tx: watch::Sender<Option<Arc<Client>>>,
     ) -> PyResult<()> {
+        let (pairing_qr_callbacks, message_callbacks, connected_callbacks, logout_callbacks) =
+            Python::attach(|py| {
+                let dispatcher = handlers.bind(py).borrow();
+                (
+                    dispatcher.pairing_qr_handlers(py),
+                    dispatcher.message_handlers(py),
+                    dispatcher.conneccted_handlers(py),
+                    dispatcher.logout_handlers(py),
+                )
+            });
+        let pairing_qr_callbacks = Arc::new(pairing_qr_callbacks);
+        let message_callbacks = Arc::new(message_callbacks);
+        let connected_callbacks = Arc::new(connected_callbacks);
+        let logout_callbacks = Arc::new(logout_callbacks);
+        info!(
+            pairing_qr_handlers = pairing_qr_callbacks.len(),
+            message_handlers = message_callbacks.len(),
+            connected_handlers = connected_callbacks.len(),
+            logout_handlers = logout_callbacks.len(),
+            "cached dispatcher handlers for runtime"
+        );
+
         info!("building WhatsApp bot");
         let mut bot = Bot::builder()
             .with_backend(backend)
             .with_transport_factory(TokioWebSocketTransportFactory::new())
             .with_http_client(UreqHttpClient::new())
             .on_event(move |event, _client| {
-                let handlers = Python::attach(|py| handlers.clone_ref(py));
                 let locals = locals.clone();
+                let pairing_qr_callbacks = Arc::clone(&pairing_qr_callbacks);
+                let message_callbacks = Arc::clone(&message_callbacks);
+                let connected_callbacks = Arc::clone(&connected_callbacks);
+                let logout_callbacks = Arc::clone(&logout_callbacks);
                 let tryx_client = Python::attach(|py| tryx_client.clone_ref(py));
                 async move {
                     match event {
                         Event::PairingQrCode { code, timeout } => {
                             info!(timeout_secs = timeout.as_secs(), "received pairing QR event");
-                            let callbacks = Python::attach(|py| {
-                                handlers
-                                    .bind(py)
-                                    .borrow()
-                                    .pairing_qr_handlers(py)
-                            });
-                            info!(handlers = callbacks.len(), "dispatching pairing QR handlers");
+                            info!(handlers = pairing_qr_callbacks.len(), "dispatching pairing QR handlers");
 
-                            for (idx, callback) in callbacks.into_iter().enumerate() {
+                            for (idx, callback) in pairing_qr_callbacks.iter().enumerate() {
                                 debug!(handler_index = idx, "calling pairing QR Python callback");
                                 let locals = locals.clone();
                                 let py_future = Python::attach(|py| -> PyResult<_> {
@@ -149,15 +217,9 @@ impl Tryx {
                         }
                         Event::Message(msg, info) => {
                             debug!(message_id = %info.id, "received message event");
-                            let callbacks = Python::attach(|py| {
-                                handlers
-                                    .bind(py)
-                                    .borrow()
-                                    .message_handlers(py)
-                            });
-                            info!(handlers = callbacks.len(), message_id = %info.id, "dispatching message handlers");
+                            info!(handlers = message_callbacks.len(), message_id = %info.id, "dispatching message handlers");
 
-                            for (idx, callback) in callbacks.into_iter().enumerate() {
+                            for (idx, callback) in message_callbacks.iter().enumerate() {
                                 debug!(handler_index = idx, message_id = %info.id, "calling message Python callback");
                                 let locals = locals.clone();
                                 let py_future = Python::attach(|py| -> PyResult<_> {
@@ -194,8 +256,7 @@ impl Tryx {
                             }
                         }
                         Event::Connected(_) => {
-                            let callbacks = Python::attach(|py| handlers.clone_ref(py).bind(py).borrow().conneccted_handlers(py));
-                            for (idx, callback) in callbacks.into_iter().enumerate() {
+                            for (idx, callback) in connected_callbacks.iter().enumerate() {
                                 debug!(handler_index = idx, "calling connected event handler");
                                 let _ = Python::attach(|py| -> PyResult<_> {
                                     let awaitable = callback.bind(py).call1((Connected{},))?;
@@ -204,6 +265,17 @@ impl Tryx {
                                 });
                             }
                         }
+                        Event::LoggedOut(logout) => {
+                            for (idx, callback) in logout_callbacks.iter().enumerate() {
+                                debug!(handler_index = idx, "calling logged out event handler");
+                                let _ = Python::attach(|py| -> PyResult<_> {
+                                    let awaitable = callback.bind(py).call1((LoggedOut::new(logout.clone()),))?;
+                                    let fut = into_future(awaitable)?;
+                                    Ok(fut)
+                                });
+                            }
+                            
+                        }
                         _ => {
                             debug!("received event without registered dispatcher path");
                         }
@@ -377,4 +449,5 @@ impl Tryx {
             })
         })
     }
+    
 }

src/dispatcher.rs

@@ -117,6 +117,14 @@ impl Dispatcher {
         debug!(handlers = handlers.len(), "collected message handlers");
         handlers
     }
+    pub fn logout_handlers(&self, py: Python<'_>) -> Vec<Py<PyAny>> {
+        let handlers = self.logged_out
+            .iter()
+            .map(|handler| handler.clone_ref(py))
+            .collect::<Vec<_>>();
+        debug!(handlers = handlers.len(), "collected logged out handlers");
+        handlers
+    }
     pub fn conneccted_handlers(&self, py: Python<'_>) -> Vec<Py<PyAny>> {
         let handlers = self.connected
             .iter()
@@ -182,7 +190,6 @@ fn dispatch_event_from_type(py: Python, event_type: &Bound<PyAny>) -> PyResult<D
         Ok(DispatchEvent::Message)
     } else {
         Err(PyErr::new::<UnsupportedEventType, _>("Unsupported event type"))
-
     }
 }
 

src/lib.rs

@@ -23,7 +23,11 @@ fn _tryx(_py: &Bound<PyModule>) -> PyResult<()> {
     let types_module = PyModule::new(_py.py(), "types")?;
     types_module.add_class::<types::JID>()?;
     types_module.add_class::<types::MessageInfo>()?;
+    types_module.add_class::<types::UploadResponse>()?;
     _py.add_submodule(&types_module)?;
+    let wacore_module = PyModule::new(_py.py(), "wacore")?;
+    wacore_module.add_class::<wacore::MediaType>()?;
+    _py.add_submodule(&wacore_module)?;
     Ok(())
 }
 
@@ -32,4 +36,5 @@ mod client;
 mod events;
 mod types;
 mod dispatcher;
-mod exceptions;
+mod exceptions;
+mod wacore;

src/types.rs

@@ -266,3 +266,19 @@ impl MessageInfo {
         format!("MessageInfo(id='{}', type='{}', push_name='{}')", self.id, self.r#type, self.push_name)
     }
 }
+
+#[pyclass]
+pub struct UploadResponse {
+    #[pyo3(get)]
+    pub url: String,
+    #[pyo3(get)]
+    pub direct_path: String,
+    #[pyo3(get)]
+    pub media_key: Vec<u8>,
+    #[pyo3(get)]
+    pub file_enc_sha256: Vec<u8>,
+    #[pyo3(get)]
+    pub file_sha256: Vec<u8>,
+    #[pyo3(get)]
+    pub file_length: u64,
+}

src/wacore.rs

@@ -0,0 +1,46 @@
+use pyo3::pyclass;
+use wacore::download::{MediaType as WACoreMediaType};
+
+#[pyclass]
+pub enum MediaType {
+    Image,
+    Video,
+    Audio,
+    Document,
+    History,
+    AppState,
+    Sticker,
+    StickerPack,
+    LinkThumbnail,
+}
+impl From<MediaType> for WACoreMediaType {
+    fn from(value: MediaType) -> Self {
+        match value {
+            MediaType::Image => WACoreMediaType::Image,
+            MediaType::Video => WACoreMediaType::Video,
+            MediaType::Audio => WACoreMediaType::Audio,
+            MediaType::Document => WACoreMediaType::Document,
+            MediaType::History => WACoreMediaType::History,
+            MediaType::AppState => WACoreMediaType::AppState,
+            MediaType::Sticker => WACoreMediaType::Sticker,
+            MediaType::StickerPack => WACoreMediaType::StickerPack,
+            MediaType::LinkThumbnail => WACoreMediaType::LinkThumbnail,
+        }
+    }
+}
+
+impl MediaType {
+    pub fn to_wacore_enum(&self) -> WACoreMediaType {
+        match &self {
+            MediaType::Image => WACoreMediaType::Image,
+            MediaType::Video => WACoreMediaType::Video,
+            MediaType::Audio => WACoreMediaType::Audio,
+            MediaType::Document => WACoreMediaType::Document,
+            MediaType::History => WACoreMediaType::History,
+            MediaType::AppState => WACoreMediaType::AppState,
+            MediaType::Sticker => WACoreMediaType::Sticker,
+            MediaType::StickerPack => WACoreMediaType::StickerPack,
+            MediaType::LinkThumbnail => WACoreMediaType::LinkThumbnail,
+        }
+    }
+}