stream from printed output

src/backend/gradio_logsview/logsview.py

@@ -1,11 +1,13 @@
 """LogsView() custom component"""
+from __future__ import annotations
 
+import io
 import logging
 import queue
 import subprocess
 import time
 import traceback
-from contextlib import contextmanager
+from contextlib import contextmanager, redirect_stdout
 from dataclasses import dataclass
 from datetime import datetime
 from functools import wraps
@@ -100,8 +102,8 @@ class LogsViewRunner:
 
         def _log(record: logging.LogRecord) -> bool:
             """Handle log record and return True if log should be yielded."""
-            if record.thread != self.thread.ident:
-                return False  # Skip if not from the thread
+            # if record.thread != self.thread.ident:
+            #     return False  # Skip if not from the thread
             if logger_name and not record.name.startswith(logger_name):
                 return False  # Skip if not from the logger
             if record.levelno < log_level:
@@ -113,17 +115,26 @@ class LogsViewRunner:
             )
             return True
 
+        stdout_queue = queue.Queue()
+        stream = LineQueueStream(stdout_queue)
         with capture_logging(log_level) as log_queue:
-            # Start thread and loop to capture and yield logs from the thread
-            self.thread.start()
-            while self.thread.is_alive():
-                while True:
-                    try:
-                        if _log(log_queue.get_nowait()):
-                            yield self.logs
-                    except queue.Empty:
-                        break
-                self.thread.join(timeout=0.1)  # adjust the timeout as needed
+            with redirect_stdout(stream):
+                # Start thread and loop to capture and yield logs from the thread
+                self.thread.start()
+                while self.thread.is_alive():
+                    while True:
+                        try:
+                            if _log(log_queue.get_nowait()):
+                                yield self.logs
+                        except queue.Empty:
+                            break
+                    while True:
+                        try:
+                            line = stdout_queue.get_nowait()
+                            yield self.log(line)
+                        except queue.Empty:
+                            break
+                    self.thread.join(timeout=0.1)  # adjust the timeout as needed
 
         # After the thread completes, yield any remaining logs
         while True:
@@ -132,6 +143,12 @@ class LogsViewRunner:
                     yield self.logs
             except queue.Empty:
                 break
+        while True:
+            try:
+                line = stdout_queue.get_nowait()
+                yield self.log(line)
+            except queue.Empty:
+                break
 
         try:
             error = error_queue.get_nowait()
@@ -287,3 +304,21 @@ def non_failing_fn(fn: Callable, *args, **kwargs) -> Callable:
             error_queue.put(e)
 
     return error_queue, _inner
+
+class LineQueueStream(io.StringIO):
+    def __init__(self, queue: queue.Queue, *args, **kwargs):
+        super().__init__(*args, **kwargs)
+        self._queue = queue
+        self._current_line = ""
+
+    def write(self, s: str) -> int:
+        if "\n" in s:
+            lines = s.split("\n")
+            for line in lines[:-1]:
+                self._current_line += line
+                self._queue.put(self._current_line)
+                self._current_line = ""
+            self._current_line += lines[-1]
+        else:
+            self._current_line += s
+        return super().write(s)