| 33 | |
| 34 | |
| 35 | class OutputBuffer: |
| 36 | LEVELS: Sequence[str] = ('info', 'warn', 'fail') |
| 37 | COLORS = {'head': 36, 'good': 32, 'warn': 33, 'fail': 31} |
| 38 | |
| 39 | # Use brighter colors on Windows for better readability. |
| 40 | if Utils.is_windows(): |
| 41 | COLORS = {'head': 96, 'good': 92, 'warn': 93, 'fail': 91} |
| 42 | |
| 43 | def __init__(self, buffer_output: bool = True) -> None: |
| 44 | self.buffer_output = buffer_output |
| 45 | self.buffer: List[str] = [] |
| 46 | self.in_section = False |
| 47 | self.section: List[str] = [] |
| 48 | self.batch = False |
| 49 | self.verbose = False |
| 50 | self.debug = False |
| 51 | self.use_colors = True |
| 52 | self.json = False |
| 53 | self.__level = 0 |
| 54 | self.__is_color_supported = ('colorama' in sys.modules) or (os.name == 'posix') |
| 55 | self.line_ended = True |
| 56 | |
| 57 | def _print(self, level: str, s: str = '', line_ended: bool = True) -> None: |
| 58 | '''Saves output to buffer (if in buffered mode), or immediately prints to stdout otherwise.''' |
| 59 | |
| 60 | # If we're logging only 'warn' or above, and this is an 'info', ignore message. |
| 61 | if self.get_level(level) < self.__level: |
| 62 | return |
| 63 | |
| 64 | if self.use_colors and self.colors_supported and len(s) > 0 and level != 'info': |
| 65 | s = "\033[0;%dm%s\033[0m" % (self.COLORS[level], s) |
| 66 | |
| 67 | if self.buffer_output: |
| 68 | # Select which list to add to. If we are in a 'with' statement, then this goes in the section buffer, otherwise the general buffer. |
| 69 | buf = self.section if self.in_section else self.buffer |
| 70 | |
| 71 | # Determine if a new line should be added, or if the last line should be appended. |
| 72 | if not self.line_ended: |
| 73 | last_entry = -1 if len(buf) > 0 else 0 |
| 74 | buf[last_entry] = buf[last_entry] + s |
| 75 | else: |
| 76 | buf.append(s) |
| 77 | |
| 78 | # When False, this tells the next call to append to the last line we just added. |
| 79 | self.line_ended = line_ended |
| 80 | else: |
| 81 | print(s) |
| 82 | |
| 83 | def get_buffer(self) -> str: |
| 84 | '''Returns all buffered output, then clears the buffer.''' |
| 85 | self.flush_section() |
| 86 | |
| 87 | buffer_str = "\n".join(self.buffer) |
| 88 | self.buffer = [] |
| 89 | return buffer_str |
| 90 | |
| 91 | def write(self) -> None: |
| 92 | '''Writes the output to stdout.''' |