Merge pull request #444 from martindurant/fix_river_async

River nodes in an async context were not returning awaitable
python-streamz · Dec 9, 2021 · 0b63090 · 0b63090
2 parents ab6c06a + f625129
commit 0b63090
Showing 1 changed file with 3 additions and 3 deletions.
diff --git a/streamz/river.py b/streamz/river.py
@@ -46,9 +46,9 @@ def update(self, x, who=None, metadata=None):
         if self.metric:
             yp = self.model.predict_one(x[0])
             weights = x[2] if len(x) > 1 else 1.0
-            self.emit(self.metric.update(x[1], yp, weights).get(), metadata=metadata)
+            return self._emit(self.metric.update(x[1], yp, weights).get(), metadata=metadata)
         if self.pass_model:
-            self.emit(self.model, metadata=metadata)
+            return self._emit(self.model, metadata=metadata)
 
 
 class RiverPredict(Stream):
@@ -59,4 +59,4 @@ def __init__(self, model, **kwargs):
 
     def update(self, x, who=None, metadata=None):
         out = self.model.predict_one(x)
-        self.emit(out, metadata=metadata)
+        return self._emit(out, metadata=metadata)