From a5a79c81941d9e0104c9af4e841a309e184caf21 Mon Sep 17 00:00:00 2001
From: Peter Schafhalter <pschafhalter@berkeley.edu>
Date: Mon, 16 Nov 2020 22:47:07 -0800
Subject: [PATCH] Implement try_read for PyExtractStream (#146)

* Implement try_read for PyExtractStream

* Fix formatting
---
 src/python/py_stream/py_extract_stream.rs | 17 ++++++++++++++++-
 1 file changed, 16 insertions(+), 1 deletion(-)

diff --git a/src/python/py_stream/py_extract_stream.rs b/src/python/py_stream/py_extract_stream.rs
index b9261554..9df38874 100644
--- a/src/python/py_stream/py_extract_stream.rs
+++ b/src/python/py_stream/py_extract_stream.rs
@@ -1,6 +1,9 @@
 use pyo3::{exceptions, prelude::*};
 
-use crate::{dataflow::stream::ExtractStream, python::PyMessage};
+use crate::{
+    dataflow::stream::{errors::TryReadError, ExtractStream},
+    python::PyMessage,
+};
 
 use super::PyReadStream;
 
@@ -43,4 +46,16 @@ impl PyExtractStream {
             ))),
         }
     }
+
+    fn try_read<'p>(&mut self) -> PyResult<Option<PyMessage>> {
+        match self.extract_stream.try_read() {
+            Ok(msg) => Ok(Some(PyMessage::from(msg))),
+            Err(TryReadError::Empty) => Ok(None),
+            Err(e) => Err(exceptions::Exception::py_err(format!(
+                "Unable to to read from stream {}: {:?}",
+                self.extract_stream.get_id(),
+                e
+            ))),
+        }
+    }
 }