Spaces:
Sleeping
Sleeping
File size: 5,077 Bytes
470be5c |
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 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 |
import streamlit as st
import io
import contextlib
import sys
import re
class _Redirect:
class IOStuff(io.StringIO):
def __init__(self, trigger, max_buffer, buffer_separator, regex, dup=None):
super().__init__()
self._trigger = trigger
self._max_buffer = max_buffer
self._buffer_separator = buffer_separator
self._regex = regex and re.compile(regex)
self._dup = dup
def write(self, __s: str) -> int:
if self._max_buffer:
concatenated_len = super().tell() + len(__s)
if concatenated_len > self._max_buffer:
rest = self.get_filtered_output()[concatenated_len - self._max_buffer:]
if self._buffer_separator is not None:
rest = rest.split(self._buffer_separator, 1)[-1]
super().seek(0)
super().write(rest)
super().truncate(super().tell() + len(__s))
res = super().write(__s)
if self._dup is not None:
self._dup.write(__s)
self._trigger(self.get_filtered_output())
return res
def get_filtered_output(self):
if self._regex is None or self._buffer_separator is None:
return self.getvalue()
return self._buffer_separator.join(filter(self._regex.search, self.getvalue().split(self._buffer_separator)))
def print_at_end(self):
self._trigger(self.get_filtered_output())
def __init__(self, stdout=None, stderr=False, format=None, to=None, max_buffer=None, buffer_separator='\n',
regex=None, duplicate_out=False):
self.io_args = {'trigger': self._write, 'max_buffer': max_buffer, 'buffer_separator': buffer_separator,
'regex': regex}
self.redirections = []
self.st = None
self.stderr = stderr is True
self.stdout = stdout is True or (stdout is None and not self.stderr)
self.format = format or 'code'
self.to = to
self.fun = None
self.duplicate_out = duplicate_out or None
self.active_nested = None
if not self.stdout and not self.stderr:
raise ValueError("one of stdout or stderr must be True")
if self.format not in ['text', 'markdown', 'latex', 'code', 'write']:
raise ValueError(
f"format need oneof the following: {', '.join(['text', 'markdown', 'latex', 'code', 'write'])}")
if self.to and (not hasattr(self.to, 'text') or not hasattr(self.to, 'empty')):
raise ValueError(f"'to' is not a streamlit container object")
def __enter__(self):
if self.st is not None:
if self.to is None:
if self.active_nested is None:
self.active_nested = self(format=self.format, max_buffer=self.io_args['max_buffer'],
buffer_separator=self.io_args['buffer_separator'],
regex=self.io_args['regex'], duplicate_out=self.duplicate_out)
return self.active_nested.__enter__()
else:
raise Exception("Already entered")
to = self.to or st
# to.text(f"{'stdout and stderr' if self.stdout and self.stderr else 'stdout' if self.stdout else 'stderr'}"
# f"{' [' + self.io_args['regex'] + ']' if self.io_args['regex'] else ''}"
# f":")
self.st = to.empty()
self.fun = getattr(self.st, self.format)
io_obj = None
def redirect(to_duplicate):
nonlocal io_obj
io_obj = _Redirect.IOStuff(dup=self.duplicate_out and to_duplicate, **self.io_args)
redirection = contextlib.redirect_stdout(io_obj)
self.redirections.append((redirection, io_obj))
redirection.__enter__()
if self.stderr:
redirect(sys.stderr)
if self.stdout:
redirect(sys.stdout)
return io_obj
def __call__(self, to=None, format=None, max_buffer=None, buffer_separator='\n', regex=None, duplicate_out=False):
return _Redirect(self.stdout, self.stderr, format=format, to=to, max_buffer=max_buffer,
buffer_separator=buffer_separator, regex=regex, duplicate_out=duplicate_out)
def __exit__(self, *exc):
if self.active_nested is not None:
nested = self.active_nested
if nested.active_nested is None:
self.active_nested = None
return nested.__exit__(*exc)
res = None
for redirection, io_obj in reversed(self.redirections):
res = redirection.__exit__(*exc)
io_obj.print_at_end()
self.redirections = []
self.st = None
self.fun = None
return res
def _write(self, data):
self.fun(data)
stdout = _Redirect()
stderr = _Redirect(stderr=True)
stdouterr = _Redirect(stdout=True, stderr=True)
|