ansi.py 8.0 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253
  1. # -*- coding: ISO-8859-1 -*-
  2. import os
  3. import re
  4. import sys
  5. terminal_escape = re.compile("(\001?\033\\[[0-9;]*m\002?)")
  6. escape_parts = re.compile("\001?\033\\[([0-9;]*)m\002?")
  7. class AnsiState(object):
  8. def __init__(
  9. self,
  10. bold=False,
  11. inverse=False,
  12. color="white",
  13. background="black",
  14. backgroundbold=False,
  15. ):
  16. self.bold = bold
  17. self.inverse = inverse
  18. self.color = color
  19. self.background = background
  20. self.backgroundbold = backgroundbold
  21. trtable = {
  22. "black": 0,
  23. "red": 4,
  24. "green": 2,
  25. "yellow": 6,
  26. "blue": 1,
  27. "magenta": 5,
  28. "cyan": 3,
  29. "white": 7,
  30. }
  31. revtable = dict(zip(trtable.values(), trtable.keys()))
  32. def get_winattr(self):
  33. attr = 0
  34. if self.bold:
  35. attr |= 0x0008
  36. if self.backgroundbold:
  37. attr |= 0x0080
  38. if self.inverse:
  39. attr |= 0x4000
  40. attr |= self.trtable[self.color]
  41. attr |= self.trtable[self.background] << 4
  42. return attr
  43. def set_winattr(self, attr):
  44. self.bold = bool(attr & 0x0008)
  45. self.backgroundbold = bool(attr & 0x0080)
  46. self.inverse = bool(attr & 0x4000)
  47. self.color = self.revtable[attr & 0x0007]
  48. self.background = self.revtable[(attr & 0x0070) >> 4]
  49. winattr = property(get_winattr, set_winattr)
  50. def __repr__(self):
  51. return (
  52. "AnsiState(bold=%s,inverse=%s,color=%9s,"
  53. "background=%9s,backgroundbold=%s)# 0x%x"
  54. % (
  55. self.bold,
  56. self.inverse,
  57. '"%s"' % self.color,
  58. '"%s"' % self.background,
  59. self.backgroundbold,
  60. self.winattr,
  61. )
  62. )
  63. def copy(self):
  64. x = AnsiState()
  65. x.bold = self.bold
  66. x.inverse = self.inverse
  67. x.color = self.color
  68. x.background = self.background
  69. x.backgroundbold = self.backgroundbold
  70. return x
  71. defaultstate = AnsiState(False, False, "white")
  72. trtable = {
  73. 0: "black",
  74. 1: "red",
  75. 2: "green",
  76. 3: "yellow",
  77. 4: "blue",
  78. 5: "magenta",
  79. 6: "cyan",
  80. 7: "white",
  81. }
  82. class AnsiWriter(object):
  83. def __init__(self, default=defaultstate):
  84. if isinstance(defaultstate, AnsiState):
  85. self.defaultstate = default
  86. else:
  87. self.defaultstate = AnsiState()
  88. self.defaultstate.winattr = defaultstate
  89. def write_color(self, text, attr=None):
  90. """write text at current cursor position and interpret color escapes.
  91. return the number of characters written.
  92. """
  93. if isinstance(attr, AnsiState):
  94. defaultstate = attr
  95. elif attr is None: # use attribute form initial console
  96. attr = self.defaultstate.copy()
  97. else:
  98. defaultstate = AnsiState()
  99. defaultstate.winattr = attr
  100. attr = defaultstate
  101. chunks = terminal_escape.split(text)
  102. n = 0 # count the characters we actually write, omitting the escapes
  103. res = []
  104. for chunk in chunks:
  105. m = escape_parts.match(chunk)
  106. if m:
  107. parts = m.group(1).split(";")
  108. if len(parts) == 1 and parts[0] == "0":
  109. attr = self.defaultstate.copy()
  110. continue
  111. for part in parts:
  112. if part == "0": # No text attribute
  113. attr = self.defaultstate.copy()
  114. attr.bold = False
  115. elif part == "7": # switch on reverse
  116. attr.inverse = True
  117. # switch on bold (i.e. intensify foreground color)
  118. elif part == "1":
  119. attr.bold = True
  120. elif len(part) == 2:
  121. if part == "22": # Normal foreground color
  122. attr.bold = False
  123. elif "30" <= part <= "37": # set foreground color
  124. attr.color = trtable[int(part) - 30]
  125. elif part == "39": # Default foreground color
  126. attr.color = self.defaultstate.color
  127. elif "40" <= part <= "47": # set background color
  128. attr.background = trtable[int(part) - 40]
  129. elif part == "49": # Default background color
  130. attr.background = self.defaultstate.background
  131. continue
  132. n += len(chunk)
  133. res.append((attr.copy(), chunk))
  134. return n, res
  135. def parse_color(self, text, attr=None):
  136. n, res = self.write_color(text, attr)
  137. return n, [attr.winattr for attr, text in res]
  138. def write_color(text, attr=None):
  139. a = AnsiWriter(defaultstate)
  140. return a.write_color(text, attr)
  141. def write_color_old(text, attr=None):
  142. """write text at current cursor position and interpret color escapes.
  143. return the number of characters written.
  144. """
  145. res = []
  146. chunks = terminal_escape.split(text)
  147. n = 0 # count the characters we actually write, omitting the escapes
  148. if attr is None: # use attribute from initial console
  149. attr = 15
  150. for chunk in chunks:
  151. m = escape_parts.match(chunk)
  152. if m:
  153. for part in m.group(1).split(";"):
  154. if part == "0": # No text attribute
  155. attr = 0
  156. elif part == "7": # switch on reverse
  157. attr |= 0x4000
  158. # switch on bold (i.e. intensify foreground color)
  159. if part == "1":
  160. attr |= 0x08
  161. elif len(part) == 2 and "30" <= part <= "37": # set foreground color
  162. part = int(part) - 30
  163. # we have to mirror bits
  164. attr = (
  165. (attr & ~0x07)
  166. | ((part & 0x1) << 2)
  167. | (part & 0x2)
  168. | ((part & 0x4) >> 2)
  169. )
  170. elif len(part) == 2 and "40" <= part <= "47": # set background color
  171. part = int(part) - 40
  172. # we have to mirror bits
  173. attr = (
  174. (attr & ~0x70)
  175. | ((part & 0x1) << 6)
  176. | ((part & 0x2) << 4)
  177. | ((part & 0x4) << 2)
  178. )
  179. # ignore blink, underline and anything we don't understand
  180. continue
  181. n += len(chunk)
  182. if chunk:
  183. res.append(("0x%x" % attr, chunk))
  184. return res
  185. # trtable={0:"black",1:"red",2:"green",3:"yellow",4:"blue",5:"magenta",6:"cyan",7:"white"}
  186. if __name__ == "__main__x":
  187. import pprint
  188. pprint = pprint.pprint
  189. s = "\033[0;31mred\033[0;32mgreen\033[0;33myellow\033[0;34mblue\033[0;35mmagenta\033[0;36mcyan\033[0;37mwhite\033[0m"
  190. pprint(write_color(s))
  191. pprint(write_color_old(s))
  192. s = "\033[1;31mred\033[1;32mgreen\033[1;33myellow\033[1;34mblue\033[1;35mmagenta\033[1;36mcyan\033[1;37mwhite\033[0m"
  193. pprint(write_color(s))
  194. pprint(write_color_old(s))
  195. s = "\033[0;7;31mred\033[0;7;32mgreen\033[0;7;33myellow\033[0;7;34mblue\033[0;7;35mmagenta\033[0;7;36mcyan\033[0;7;37mwhite\033[0m"
  196. pprint(write_color(s))
  197. pprint(write_color_old(s))
  198. s = "\033[1;7;31mred\033[1;7;32mgreen\033[1;7;33myellow\033[1;7;34mblue\033[1;7;35mmagenta\033[1;7;36mcyan\033[1;7;37mwhite\033[0m"
  199. pprint(write_color(s))
  200. pprint(write_color_old(s))
  201. if __name__ == "__main__":
  202. import pprint
  203. import console
  204. pprint = pprint.pprint
  205. c = console.Console()
  206. c.write_color("dhsjdhs")
  207. c.write_color("\033[0;32mIn [\033[1;32m1\033[0;32m]:")
  208. print
  209. pprint(write_color("\033[0;32mIn [\033[1;32m1\033[0;32m]:"))
  210. if __name__ == "__main__x":
  211. import pprint
  212. pprint = pprint.pprint
  213. s = "\033[0;31mred\033[0;32mgreen\033[0;33myellow\033[0;34mblue\033[0;35mmagenta\033[0;36mcyan\033[0;37mwhite\033[0m"
  214. pprint(write_color(s))