Hide keyboard shortcuts

Hot-keys on this page

r m x p   toggle line displays

j k   next/prev highlighted chunk

0   (zero) top of page

1   (one) first highlighted chunk

1"""Helper functions for writing to terminals and files.""" 

2import os 

3import shutil 

4import sys 

5from typing import Optional 

6from typing import Sequence 

7from typing import TextIO 

8 

9from .wcwidth import wcswidth 

10 

11 

12# This code was initially copied from py 1.8.1, file _io/terminalwriter.py. 

13 

14 

15def get_terminal_width() -> int: 

16 width, _ = shutil.get_terminal_size(fallback=(80, 24)) 

17 

18 # The Windows get_terminal_size may be bogus, let's sanify a bit. 

19 if width < 40: 

20 width = 80 

21 

22 return width 

23 

24 

25def should_do_markup(file: TextIO) -> bool: 

26 if os.environ.get("PY_COLORS") == "1": 

27 return True 

28 if os.environ.get("PY_COLORS") == "0": 

29 return False 

30 if "NO_COLOR" in os.environ: 

31 return False 

32 if "FORCE_COLOR" in os.environ: 

33 return True 

34 return ( 

35 hasattr(file, "isatty") and file.isatty() and os.environ.get("TERM") != "dumb" 

36 ) 

37 

38 

39class TerminalWriter: 

40 _esctable = dict( 

41 black=30, 

42 red=31, 

43 green=32, 

44 yellow=33, 

45 blue=34, 

46 purple=35, 

47 cyan=36, 

48 white=37, 

49 Black=40, 

50 Red=41, 

51 Green=42, 

52 Yellow=43, 

53 Blue=44, 

54 Purple=45, 

55 Cyan=46, 

56 White=47, 

57 bold=1, 

58 light=2, 

59 blink=5, 

60 invert=7, 

61 ) 

62 

63 def __init__(self, file: Optional[TextIO] = None) -> None: 

64 if file is None: 

65 file = sys.stdout 

66 if hasattr(file, "isatty") and file.isatty() and sys.platform == "win32": 

67 try: 

68 import colorama 

69 except ImportError: 

70 pass 

71 else: 

72 file = colorama.AnsiToWin32(file).stream 

73 assert file is not None 

74 self._file = file 

75 self.hasmarkup = should_do_markup(file) 

76 self._current_line = "" 

77 self._terminal_width = None # type: Optional[int] 

78 self.code_highlight = True 

79 

80 @property 

81 def fullwidth(self) -> int: 

82 if self._terminal_width is not None: 

83 return self._terminal_width 

84 return get_terminal_width() 

85 

86 @fullwidth.setter 

87 def fullwidth(self, value: int) -> None: 

88 self._terminal_width = value 

89 

90 @property 

91 def width_of_current_line(self) -> int: 

92 """Return an estimate of the width so far in the current line.""" 

93 return wcswidth(self._current_line) 

94 

95 def markup(self, text: str, **markup: bool) -> str: 

96 for name in markup: 

97 if name not in self._esctable: 

98 raise ValueError("unknown markup: {!r}".format(name)) 

99 if self.hasmarkup: 

100 esc = [self._esctable[name] for name, on in markup.items() if on] 

101 if esc: 

102 text = "".join("\x1b[%sm" % cod for cod in esc) + text + "\x1b[0m" 

103 return text 

104 

105 def sep( 

106 self, 

107 sepchar: str, 

108 title: Optional[str] = None, 

109 fullwidth: Optional[int] = None, 

110 **markup: bool 

111 ) -> None: 

112 if fullwidth is None: 

113 fullwidth = self.fullwidth 

114 # the goal is to have the line be as long as possible 

115 # under the condition that len(line) <= fullwidth 

116 if sys.platform == "win32": 

117 # if we print in the last column on windows we are on a 

118 # new line but there is no way to verify/neutralize this 

119 # (we may not know the exact line width) 

120 # so let's be defensive to avoid empty lines in the output 

121 fullwidth -= 1 

122 if title is not None: 

123 # we want 2 + 2*len(fill) + len(title) <= fullwidth 

124 # i.e. 2 + 2*len(sepchar)*N + len(title) <= fullwidth 

125 # 2*len(sepchar)*N <= fullwidth - len(title) - 2 

126 # N <= (fullwidth - len(title) - 2) // (2*len(sepchar)) 

127 N = max((fullwidth - len(title) - 2) // (2 * len(sepchar)), 1) 

128 fill = sepchar * N 

129 line = "{} {} {}".format(fill, title, fill) 

130 else: 

131 # we want len(sepchar)*N <= fullwidth 

132 # i.e. N <= fullwidth // len(sepchar) 

133 line = sepchar * (fullwidth // len(sepchar)) 

134 # in some situations there is room for an extra sepchar at the right, 

135 # in particular if we consider that with a sepchar like "_ " the 

136 # trailing space is not important at the end of the line 

137 if len(line) + len(sepchar.rstrip()) <= fullwidth: 

138 line += sepchar.rstrip() 

139 

140 self.line(line, **markup) 

141 

142 def write(self, msg: str, *, flush: bool = False, **markup: bool) -> None: 

143 if msg: 

144 current_line = msg.rsplit("\n", 1)[-1] 

145 if "\n" in msg: 

146 self._current_line = current_line 

147 else: 

148 self._current_line += current_line 

149 

150 msg = self.markup(msg, **markup) 

151 

152 try: 

153 self._file.write(msg) 

154 except UnicodeEncodeError: 

155 # Some environments don't support printing general Unicode 

156 # strings, due to misconfiguration or otherwise; in that case, 

157 # print the string escaped to ASCII. 

158 # When the Unicode situation improves we should consider 

159 # letting the error propagate instead of masking it (see #7475 

160 # for one brief attempt). 

161 msg = msg.encode("unicode-escape").decode("ascii") 

162 self._file.write(msg) 

163 

164 if flush: 

165 self.flush() 

166 

167 def line(self, s: str = "", **markup: bool) -> None: 

168 self.write(s, **markup) 

169 self.write("\n") 

170 

171 def flush(self) -> None: 

172 self._file.flush() 

173 

174 def _write_source(self, lines: Sequence[str], indents: Sequence[str] = ()) -> None: 

175 """Write lines of source code possibly highlighted. 

176 

177 Keeping this private for now because the API is clunky. We should discuss how 

178 to evolve the terminal writer so we can have more precise color support, for example 

179 being able to write part of a line in one color and the rest in another, and so on. 

180 """ 

181 if indents and len(indents) != len(lines): 

182 raise ValueError( 

183 "indents size ({}) should have same size as lines ({})".format( 

184 len(indents), len(lines) 

185 ) 

186 ) 

187 if not indents: 

188 indents = [""] * len(lines) 

189 source = "\n".join(lines) 

190 new_lines = self._highlight(source).splitlines() 

191 for indent, new_line in zip(indents, new_lines): 

192 self.line(indent + new_line) 

193 

194 def _highlight(self, source: str) -> str: 

195 """Highlight the given source code if we have markup support.""" 

196 if not self.hasmarkup or not self.code_highlight: 

197 return source 

198 try: 

199 from pygments.formatters.terminal import TerminalFormatter 

200 from pygments.lexers.python import PythonLexer 

201 from pygments import highlight 

202 except ImportError: 

203 return source 

204 else: 

205 highlighted = highlight( 

206 source, PythonLexer(), TerminalFormatter(bg="dark") 

207 ) # type: str 

208 return highlighted