Spaces:
Sleeping
Sleeping
File size: 3,751 Bytes
1df74c6 |
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 |
#!/usr/bin/env python3
import re
import sys
from contextlib import contextmanager
class ANSI:
ansi_color = {
"black": "\033[30m",
"red": "\033[31m",
"green": "\033[32m",
"yellow": "\033[33m",
"blue": "\033[34m",
"purple": "\033[35m",
"blue_light": "\033[36m",
"white": "\033[37m",
"reset": "\033[0m",
"upline": "\033[1A",
"clear_line": "\033[2K",
"clear": "\033[2J",
}
ansi_nocolor = {
"black": "",
"red": "",
"green": "",
"yellow": "",
"blue": "",
"purple": "",
"blue_light": "",
"white": "",
"reset": "",
"upline": "\033[1A\033[",
"clear_line": "\033[K",
"clear": "\033[2J",
}
def __init__(self):
self._dict = ANSI.ansi_color if ("--color" in sys.argv) else ANSI.ansi_nocolor
def switch(self, color: bool):
self._dict = ANSI.ansi_color if color else ANSI.ansi_nocolor
def keys(self):
return self._dict.keys()
def items(self):
return self._dict.items()
def __getitem__(self, key):
return self._dict[key]
def __str__(self):
return str(self._dict)
def __repr__(self):
return repr(self._dict)
ansi = ANSI()
def remove_ansi(s: str) -> str:
ansi_escape = re.compile(r"(\x9B|\x1B\[)[0-?]*[ -\/]*[@-~]")
return ansi_escape.sub("", s)
def get_ansi_len(s: str) -> int:
return len(s) - len(remove_ansi(s))
def prints(*args: str, indent: int = 0, prefix: str = "", **kwargs):
assert indent >= 0
new_args = []
for arg in args:
new_args.append(indent_str(str(arg), indent=indent))
if len(new_args):
new_args[0] = prefix + str(new_args[0])
print(*new_args, **kwargs)
def output_iter(_iter: int, iteration: int = None, iter_len: int = 4) -> str:
if iteration is None:
pattern = "{blue_light}[ {red}{0}{blue_light} ]{reset}"
return pattern.format(str(_iter).rjust(iter_len), **ansi)
else:
iter_str = str(iteration)
length = len(iter_str)
pattern = (
"{blue_light}[ {red}{0}{blue_light} " "/ {red}{1}{blue_light} ]{reset}"
)
return pattern.format(str(_iter).rjust(length), iter_str, **ansi)
def indent_str(s_: str, indent: int = 0) -> str:
# modified from torch.nn.modules._addindent
if indent > 0 and s_:
s_ = indent * " " + str(s_[:-1]).replace("\n", "\n" + indent * " ") + s_[-1]
return s_
class IndentRedirect: # TODO: inherit TextIOWrapper?
def __init__(self, buffer: bool = True, indent: int = 0):
self.__console__ = sys.stdout
self.indent = indent
self.__buffer: str = None
if buffer:
self.__buffer = ""
def write(self, text: str, indent: int = None):
indent = indent if indent is not None else self.indent
text = indent_str(text, indent=indent)
if self.__buffer is None:
self.__console__.write(text)
else:
self.__buffer += text
def flush(self):
if self.__buffer is not None:
self.__console__.write(self.__buffer)
self.__buffer = ""
self.__console__.flush()
@contextmanager
def __call__(self) -> None:
try:
sys.stdout = self
yield
finally:
sys.stdout = self.__console__
self.__buffer = ""
def enable(self):
sys.stdout = self
def disable(self):
if self.__buffer is not None:
self.__buffer = ""
sys.stdout = self.__console__
@property
def buffer(self) -> str:
return self.__buffer
redirect = IndentRedirect()
|