File size: 5,628 Bytes
91aba60
612f625
91aba60
612f625
 
 
91aba60
 
 
 
 
 
 
612f625
91aba60
 
 
 
 
 
 
 
 
 
 
 
 
612f625
91aba60
 
 
 
 
 
 
 
 
 
 
 
 
 
 
612f625
 
 
91aba60
 
 
 
612f625
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
91aba60
 
 
 
612f625
91aba60
 
 
 
 
 
 
 
612f625
91aba60
612f625
 
91aba60
612f625
ef51884
91aba60
 
 
 
 
612f625
 
 
 
 
 
 
91aba60
 
 
 
 
612f625
 
 
 
 
 
91aba60
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
612f625
 
 
 
 
 
 
 
 
 
 
91aba60
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
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
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
import contextlib
import io
import re
import sys

import streamlit as st


class _Redirect:
    """
    Based on: https://gist.github.com/schaumb/037f139035d93cff3ad9f4f7e5f739ce
    Also see: https://github.com/streamlit/streamlit/issues/268#issuecomment-810478208
    """

    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: {self.to}")

    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"Redirected output from "
            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(max_buffer=1, buffer_separator="\n")
stderr = _Redirect(stderr=True)
stdouterr = _Redirect(stdout=True, stderr=True)