yt_dlp/minicurses.py

   1 import functools
   2 from threading import Lock
   3 from .utils import supports_terminal_sequences, TERMINAL_SEQUENCES, write_string
   4
   5
   6 class MultilinePrinterBase:
   7     def __init__(self, stream=None, lines=1):
   8         self.stream = stream
   9         self.maximum = lines - 1
  10
  11     def __enter__(self):
  12         return self
  13
  14     def __exit__(self, *args):
  15         self.end()
  16
  17     def print_at_line(self, text, pos):
  18         pass
  19
  20     def end(self):
  21         pass
  22
  23     def _add_line_number(self, text, line):
  24         if self.maximum:
  25             return f'{line + 1}: {text}'
  26         return text
  27
  28     def write(self, *text):
  29         write_string(''.join(text), self.stream)
  30
  31
  32 class QuietMultilinePrinter(MultilinePrinterBase):
  33     pass
  34
  35
  36 class MultilineLogger(MultilinePrinterBase):
  37     def write(self, *text):
  38         self.stream.debug(''.join(text))
  39
  40     def print_at_line(self, text, pos):
  41         # stream is the logger object, not an actual stream
  42         self.write(self._add_line_number(text, pos))
  43
  44
  45 class BreaklineStatusPrinter(MultilinePrinterBase):
  46     def print_at_line(self, text, pos):
  47         self.write(self._add_line_number(text, pos), '\n')
  48
  49
  50 class MultilinePrinter(MultilinePrinterBase):
  51     def __init__(self, stream=None, lines=1, preserve_output=True):
  52         super().__init__(stream, lines)
  53         self.preserve_output = preserve_output
  54         self._lastline = self._lastlength = 0
  55         self._movelock = Lock()
  56         self._HAVE_FULLCAP = supports_terminal_sequences(self.stream)
  57
  58     def lock(func):
  59         @functools.wraps(func)
  60         def wrapper(self, *args, **kwargs):
  61             with self._movelock:
  62                 return func(self, *args, **kwargs)
  63         return wrapper
  64
  65     def _move_cursor(self, dest):
  66         current = min(self._lastline, self.maximum)
  67         yield '\r'
  68         distance = dest - current
  69         if distance < 0:
  70             yield TERMINAL_SEQUENCES['UP'] * -distance
  71         elif distance > 0:
  72             yield TERMINAL_SEQUENCES['DOWN'] * distance
  73         self._lastline = dest
  74
  75     @lock
  76     def print_at_line(self, text, pos):
  77         if self._HAVE_FULLCAP:
  78             self.write(*self._move_cursor(pos), TERMINAL_SEQUENCES['ERASE_LINE'], text)
  79
  80         text = self._add_line_number(text, pos)
  81         textlen = len(text)
  82         if self._lastline == pos:
  83             # move cursor at the start of progress when writing to same line
  84             prefix = '\r'
  85             if self._lastlength > textlen:
  86                 text += ' ' * (self._lastlength - textlen)
  87             self._lastlength = textlen
  88         else:
  89             # otherwise, break the line
  90             prefix = '\n'
  91             self._lastlength = textlen
  92         self.write(prefix, text)
  93         self._lastline = pos
  94
  95     @lock
  96     def end(self):
  97         # move cursor to the end of the last line, and write line break
  98         # so that other to_screen calls can precede
  99         text = self._move_cursor(self.maximum) if self._HAVE_FULLCAP else []
 100         if self.preserve_output:
 101             self.write(*text, '\n')
 102             return
 103
 104         if self._HAVE_FULLCAP:
 105             self.write(
 106                 *text, TERMINAL_SEQUENCES['ERASE_LINE'],
 107                 f'{TERMINAL_SEQUENCES["UP"]}{TERMINAL_SEQUENCES["ERASE_LINE"]}' * self.maximum)
 108         else:
 109             self.write(*text, ' ' * self._lastlength)