yt_dlp/minicurses.py

   1 import functools
   2 from threading import Lock
   3 from .utils import supports_terminal_sequences, write_string
   4
   5
   6 CONTROL_SEQUENCES = {
   7     'DOWN': '\n',
   8     'UP': '\033[A',
   9     'ERASE_LINE': '\033[K',
  10     'RESET': '\033[0m',
  11 }
  12
  13
  14 _COLORS = {
  15     'BLACK': '0',
  16     'RED': '1',
  17     'GREEN': '2',
  18     'YELLOW': '3',
  19     'BLUE': '4',
  20     'PURPLE': '5',
  21     'CYAN': '6',
  22     'WHITE': '7',
  23 }
  24
  25
  26 _TEXT_STYLES = {
  27     'NORMAL': '0',
  28     'BOLD': '1',
  29     'UNDERLINED': '4',
  30 }
  31
  32
  33 def format_text(text, f):
  34     '''
  35     @param f    String representation of formatting to apply in the form:
  36                 [style] [light] font_color [on [light] bg_color]
  37                 Eg: "red", "bold green on light blue"
  38     '''
  39     f = f.upper()
  40     tokens = f.strip().split()
  41
  42     bg_color = ''
  43     if 'ON' in tokens:
  44         if tokens[-1] == 'ON':
  45             raise SyntaxError(f'Empty background format specified in {f!r}')
  46         if tokens[-1] not in _COLORS:
  47             raise SyntaxError(f'{tokens[-1]} in {f!r} must be a color')
  48         bg_color = f'4{_COLORS[tokens.pop()]}'
  49         if tokens[-1] == 'LIGHT':
  50             bg_color = f'0;10{bg_color[1:]}'
  51             tokens.pop()
  52         if tokens[-1] != 'ON':
  53             raise SyntaxError(f'Invalid format {f.split(" ON ", 1)[1]!r} in {f!r}')
  54         bg_color = f'\033[{bg_color}m'
  55         tokens.pop()
  56
  57     if not tokens:
  58         fg_color = ''
  59     elif tokens[-1] not in _COLORS:
  60         raise SyntaxError(f'{tokens[-1]} in {f!r} must be a color')
  61     else:
  62         fg_color = f'3{_COLORS[tokens.pop()]}'
  63         if tokens and tokens[-1] == 'LIGHT':
  64             fg_color = f'9{fg_color[1:]}'
  65             tokens.pop()
  66         fg_style = tokens.pop() if tokens and tokens[-1] in _TEXT_STYLES else 'NORMAL'
  67         fg_color = f'\033[{_TEXT_STYLES[fg_style]};{fg_color}m'
  68         if tokens:
  69             raise SyntaxError(f'Invalid format {" ".join(tokens)!r} in {f!r}')
  70
  71     if fg_color or bg_color:
  72         return f'{fg_color}{bg_color}{text}{CONTROL_SEQUENCES["RESET"]}'
  73     else:
  74         return text
  75
  76
  77 class MultilinePrinterBase:
  78     def __init__(self, stream=None, lines=1):
  79         self.stream = stream
  80         self.maximum = lines - 1
  81
  82     def __enter__(self):
  83         return self
  84
  85     def __exit__(self, *args):
  86         self.end()
  87
  88     def print_at_line(self, text, pos):
  89         pass
  90
  91     def end(self):
  92         pass
  93
  94     def _add_line_number(self, text, line):
  95         if self.maximum:
  96             return f'{line + 1}: {text}'
  97         return text
  98
  99     def write(self, *text):
 100         write_string(''.join(text), self.stream)
 101
 102
 103 class QuietMultilinePrinter(MultilinePrinterBase):
 104     pass
 105
 106
 107 class MultilineLogger(MultilinePrinterBase):
 108     def write(self, *text):
 109         self.stream.debug(''.join(text))
 110
 111     def print_at_line(self, text, pos):
 112         # stream is the logger object, not an actual stream
 113         self.write(self._add_line_number(text, pos))
 114
 115
 116 class BreaklineStatusPrinter(MultilinePrinterBase):
 117     def print_at_line(self, text, pos):
 118         self.write(self._add_line_number(text, pos), '\n')
 119
 120
 121 class MultilinePrinter(MultilinePrinterBase):
 122     def __init__(self, stream=None, lines=1, preserve_output=True):
 123         super().__init__(stream, lines)
 124         self.preserve_output = preserve_output
 125         self._lastline = self._lastlength = 0
 126         self._movelock = Lock()
 127         self._HAVE_FULLCAP = supports_terminal_sequences(self.stream)
 128
 129     def lock(func):
 130         @functools.wraps(func)
 131         def wrapper(self, *args, **kwargs):
 132             with self._movelock:
 133                 return func(self, *args, **kwargs)
 134         return wrapper
 135
 136     def _move_cursor(self, dest):
 137         current = min(self._lastline, self.maximum)
 138         yield '\r'
 139         distance = dest - current
 140         if distance < 0:
 141             yield CONTROL_SEQUENCES['UP'] * -distance
 142         elif distance > 0:
 143             yield CONTROL_SEQUENCES['DOWN'] * distance
 144         self._lastline = dest
 145
 146     @lock
 147     def print_at_line(self, text, pos):
 148         if self._HAVE_FULLCAP:
 149             self.write(*self._move_cursor(pos), CONTROL_SEQUENCES['ERASE_LINE'], text)
 150
 151         text = self._add_line_number(text, pos)
 152         textlen = len(text)
 153         if self._lastline == pos:
 154             # move cursor at the start of progress when writing to same line
 155             prefix = '\r'
 156             if self._lastlength > textlen:
 157                 text += ' ' * (self._lastlength - textlen)
 158             self._lastlength = textlen
 159         else:
 160             # otherwise, break the line
 161             prefix = '\n'
 162             self._lastlength = textlen
 163         self.write(prefix, text)
 164         self._lastline = pos
 165
 166     @lock
 167     def end(self):
 168         # move cursor to the end of the last line, and write line break
 169         # so that other to_screen calls can precede
 170         text = self._move_cursor(self.maximum) if self._HAVE_FULLCAP else []
 171         if self.preserve_output:
 172             self.write(*text, '\n')
 173             return
 174
 175         if self._HAVE_FULLCAP:
 176             self.write(
 177                 *text, CONTROL_SEQUENCES['ERASE_LINE'],
 178                 f'{CONTROL_SEQUENCES["UP"]}{CONTROL_SEQUENCES["ERASE_LINE"]}' * self.maximum)
 179         else:
 180             self.write(*text, ' ' * self._lastlength)