From dbdbd2d4c15560b52a0ef322ba34a46f57d8703a Mon Sep 17 00:00:00 2001 From: Shroominic Date: Fri, 1 Dec 2023 20:12:47 +0100 Subject: [PATCH] =?UTF-8?q?=F0=9F=94=96=20support=20rich=20stream=20output?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- src/funcchain/streaming.py | 5 ++++- 1 file changed, 4 insertions(+), 1 deletion(-) diff --git a/src/funcchain/streaming.py b/src/funcchain/streaming.py index 163a9c8..34b3c68 100644 --- a/src/funcchain/streaming.py +++ b/src/funcchain/streaming.py @@ -85,9 +85,12 @@ def stream_to( >>> with stream_to(print): ... # your chain calls here """ + import builtins, rich + # TODO: set streaming to true specifically in this context - if fn is print and kwargs == {}: + if (fn is builtins.print or fn is rich.print) and kwargs == {}: kwargs = {"end": "", "flush": True} + cb = AsyncStreamHandler(fn, kwargs) stream_handler.set(cb) yield cb