]>
Commit | Line | Data |
---|---|---|
bd50a52b THD |
1 | import os |
2 | ||
3 | from threading import Lock | |
4 | from .utils import compat_os_name, get_windows_version | |
5 | ||
6 | ||
7 | class MultilinePrinterBase(): | |
8 | def __enter__(self): | |
9 | return self | |
10 | ||
11 | def __exit__(self, *args): | |
12 | self.end() | |
13 | ||
14 | def print_at_line(self, text, pos): | |
15 | pass | |
16 | ||
17 | def end(self): | |
18 | pass | |
19 | ||
20 | ||
21 | class MultilinePrinter(MultilinePrinterBase): | |
22 | ||
23 | def __init__(self, stream, lines): | |
24 | """ | |
25 | @param stream stream to write to | |
26 | @lines number of lines to be written | |
27 | """ | |
28 | self.stream = stream | |
29 | ||
30 | is_win10 = compat_os_name == 'nt' and get_windows_version() >= (10, ) | |
31 | self.CARRIAGE_RETURN = '\r' | |
32 | if os.getenv('TERM') and self._isatty() or is_win10: | |
33 | # reason not to use curses https://github.com/yt-dlp/yt-dlp/pull/1036#discussion_r713851492 | |
34 | # escape sequences for Win10 https://docs.microsoft.com/en-us/windows/console/console-virtual-terminal-sequences | |
35 | self.UP = '\x1b[A' | |
36 | self.DOWN = '\n' | |
37 | self.ERASE_LINE = '\x1b[K' | |
38 | self._HAVE_FULLCAP = self._isatty() or is_win10 | |
39 | else: | |
40 | self.UP = self.DOWN = self.ERASE_LINE = None | |
41 | self._HAVE_FULLCAP = False | |
42 | ||
43 | # lines are numbered from top to bottom, counting from 0 to self.maximum | |
44 | self.maximum = lines - 1 | |
45 | self.lastline = 0 | |
46 | self.lastlength = 0 | |
47 | ||
48 | self.movelock = Lock() | |
49 | ||
50 | @property | |
51 | def have_fullcap(self): | |
52 | """ | |
53 | True if the TTY is allowing to control cursor, | |
54 | so that multiline progress works | |
55 | """ | |
56 | return self._HAVE_FULLCAP | |
57 | ||
58 | def _isatty(self): | |
59 | try: | |
60 | return self.stream.isatty() | |
61 | except BaseException: | |
62 | return False | |
63 | ||
64 | def _move_cursor(self, dest): | |
65 | current = min(self.lastline, self.maximum) | |
66 | self.stream.write(self.CARRIAGE_RETURN) | |
67 | if current == dest: | |
68 | # current and dest are at same position, no need to move cursor | |
69 | return | |
70 | elif current > dest: | |
71 | # when maximum == 2, | |
72 | # 0. dest | |
73 | # 1. | |
74 | # 2. current | |
75 | self.stream.write(self.UP * (current - dest)) | |
76 | elif current < dest: | |
77 | # when maximum == 2, | |
78 | # 0. current | |
79 | # 1. | |
80 | # 2. dest | |
81 | self.stream.write(self.DOWN * (dest - current)) | |
82 | self.lastline = dest | |
83 | ||
84 | def print_at_line(self, text, pos): | |
85 | with self.movelock: | |
86 | if self.have_fullcap: | |
87 | self._move_cursor(pos) | |
88 | self.stream.write(self.ERASE_LINE) | |
89 | self.stream.write(text) | |
90 | else: | |
91 | if self.maximum != 0: | |
92 | # let user know about which line is updating the status | |
93 | text = f'{pos + 1}: {text}' | |
94 | textlen = len(text) | |
95 | if self.lastline == pos: | |
96 | # move cursor at the start of progress when writing to same line | |
97 | self.stream.write(self.CARRIAGE_RETURN) | |
98 | if self.lastlength > textlen: | |
99 | text += ' ' * (self.lastlength - textlen) | |
100 | self.lastlength = textlen | |
101 | else: | |
102 | # otherwise, break the line | |
103 | self.stream.write('\n') | |
104 | self.lastlength = 0 | |
105 | self.stream.write(text) | |
106 | self.lastline = pos | |
107 | ||
108 | def end(self): | |
109 | with self.movelock: | |
110 | # move cursor to the end of the last line, and write line break | |
111 | # so that other to_screen calls can precede | |
112 | self._move_cursor(self.maximum) | |
113 | self.stream.write('\n') | |
114 | ||
115 | ||
116 | class QuietMultilinePrinter(MultilinePrinterBase): | |
117 | def __init__(self): | |
118 | self.have_fullcap = True | |
119 | ||
120 | ||
121 | class BreaklineStatusPrinter(MultilinePrinterBase): | |
122 | ||
123 | def __init__(self, stream, lines): | |
124 | """ | |
125 | @param stream stream to write to | |
126 | """ | |
127 | self.stream = stream | |
128 | self.maximum = lines | |
129 | self.have_fullcap = True | |
130 | ||
131 | def print_at_line(self, text, pos): | |
132 | if self.maximum != 0: | |
133 | # let user know about which line is updating the status | |
134 | text = f'{pos + 1}: {text}' | |
135 | self.stream.write(text + '\n') |