yt_dlp/minicurses.py

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