Move message parsing to PBClient

1 files changed, 23 insertions, 0 deletions

diff --git a/src/main.rs b/src/main.rs
index ea58059..b12f5a6 100644
--- a/src/main.rs
+++ b/src/main.rs
@@ -15,6 +15,7 @@ use websocket::Receiver;
 use websocket::WebSocketStream;
 use websocket::DataFrame;
 use websocket::Client;
+
 struct PBClient {
     client: websocket::client::Client<websocket::dataframe::DataFrame,
                                       websocket::client::sender::Sender<websocket::stream::WebSocketStream>,
@@ -36,6 +37,28 @@ impl PBClient {
             client: response.begin()
         }
     }
+
+    fn process_message(message: Result<websocket::Message, websocket::result::WebSocketError>) -> Option<Message> {
+        let message = match message {
+            Ok(m) => m,
+            Err(e) => {
+                println!("Error: {:?}", e);
+                return None;
+            }
+        };
+
+        if let websocket::Message::Text(message) = message {
+                Message::parse(message.as_ref())
+        }
+        else {
+            None
+        }
+    }
+
+    fn messages<'a>(&'a mut self) -> Box<Iterator<Item=Message> + 'a> {
+        let mut receiver = self.client.get_mut_reciever(); // there is a typo in the API
+        Box::new(receiver.incoming_messages().filter_map(PBClient::process_message))
+    }
 }
 
 fn main() {