File size: 1,686 Bytes
448ed98 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 |
import streamlit as st
import io
from contextlib import contextmanager
import contextlib
import re
@contextmanager
def st_capture(output_func):
with io.StringIO() as stdout, contextlib.redirect_stdout(stdout):
old_write = stdout.write
def new_write(string):
ret = old_write(string)
output_func(stdout.getvalue())
return ret
stdout.write = new_write
yield
class TqdmToStreamlit(io.StringIO):
"""
A custom file-like object that redirects tqdm's output to Streamlit's st.progress
and st.text widgets.
"""
# def __init__(self, progress_bar, text_element):
# super().__init__()
# self.progress_bar = progress_bar
# self.text_element = text_element
# self.last_progress = 0
def __init__(self, text_element):
super().__init__()
self.text_element = text_element
def write(self, buf):
# We can use this code if we want to have a nicer progress bar.
# match = re.search(r"(\d+)%\|", buf)
# if match:
# progress = int(match.group(1))
# if progress > self.last_progress:
# self.progress_bar.progress(progress)
# self.text_element.text(buf.strip())
# self.last_progress = progress
self.text_element.text(buf.strip())
def flush(self):
pass # No-op
class QueueIO(io.StringIO):
"""
A custom file-like object that writes to a queue.
Used to capture stdout/stderr.
"""
def __init__(self, q):
self.queue = q
def write(self, buf):
self.queue.put(buf)
def flush(self):
pass |