File size: 804 Bytes
37f86d1
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
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
"""
`sys.stdout` wrapper for both Google Colab and local environment.
"""

import sys
import tempfile


class StdoutWrapper:
    def __init__(self):
        self.temp_file = tempfile.NamedTemporaryFile(mode="w+", delete=False)
        self.original_stdout = sys.stdout

    def write(self, message: str):
        self.temp_file.write(message)
        self.temp_file.flush()
        print(message, end="", file=self.original_stdout)

    def flush(self):
        self.temp_file.flush()

    def read(self):
        self.temp_file.seek(0)
        return self.temp_file.read()

    def close(self):
        self.temp_file.close()

    def fileno(self):
        return self.temp_file.fileno()


try:
    import google.colab

    SAFE_STDOUT = StdoutWrapper()
except ImportError:
    SAFE_STDOUT = sys.stdout