terminal.py 4.6 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129
  1. #!/usr/bin/env python3
  2. # Contest Management System - http://cms-dev.github.io/
  3. # Copyright © 2015 Luca Versari <veluca93@gmail.com>
  4. # Copyright © 2018 Luca Chiodini <luca@chiodini.org>
  5. #
  6. # This program is free software: you can redistribute it and/or modify
  7. # it under the terms of the GNU Affero General Public License as
  8. # published by the Free Software Foundation, either version 3 of the
  9. # License, or (at your option) any later version.
  10. #
  11. # This program is distributed in the hope that it will be useful,
  12. # but WITHOUT ANY WARRANTY; without even the implied warranty of
  13. # MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
  14. # GNU Affero General Public License for more details.
  15. #
  16. # You should have received a copy of the GNU Affero General Public License
  17. # along with this program. If not, see <http://www.gnu.org/licenses/>.
  18. import curses
  19. import sys
  20. class colors:
  21. BLACK = curses.COLOR_BLACK
  22. RED = curses.COLOR_RED
  23. GREEN = curses.COLOR_GREEN
  24. YELLOW = curses.COLOR_YELLOW
  25. BLUE = curses.COLOR_BLUE
  26. MAGENTA = curses.COLOR_MAGENTA
  27. CYAN = curses.COLOR_CYAN
  28. WHITE = curses.COLOR_WHITE
  29. class directions:
  30. UP = 1
  31. DOWN = 2
  32. LEFT = 3
  33. RIGHT = 4
  34. def has_color_support(stream):
  35. """Try to determine if the given stream supports colored output.
  36. Return True only if the stream declares to be a TTY, if it has a
  37. file descriptor on which ncurses can initialize a terminal and if
  38. that terminal's entry in terminfo declares support for colors.
  39. stream (fileobj): a file-like object (that adheres to the API
  40. declared in the `io' package).
  41. return (bool): True if we're sure that colors are supported, False
  42. if they aren't or if we can't tell.
  43. """
  44. if stream.isatty():
  45. try:
  46. curses.setupterm(fd=stream.fileno())
  47. # See `man terminfo` for capabilities' names and meanings.
  48. if curses.tigetnum("colors") > 0:
  49. return True
  50. # fileno() can raise OSError.
  51. except Exception:
  52. pass
  53. return False
  54. def add_color_to_string(string, color, stream=sys.stdout, bold=False,
  55. force=False):
  56. """Format the string to be printed with the given color.
  57. Insert formatting characters that, when printed on a terminal, will
  58. make the given string appear with the given foreground color if the
  59. stream passed has color support. Else return the string as it is.
  60. string (string): the string to color.
  61. color (int): the color as a colors constant, like colors.BLACK.
  62. stream (fileobj): a file-like object (that adheres to the API
  63. declared in the `io' package). Defaults to sys.stdout.
  64. bold (bool): True if the string should be bold.
  65. force (bool): True if the string should be formatted even if the
  66. given stream has no color support.
  67. return (string): the formatted string.
  68. """
  69. if force or has_color_support(stream):
  70. return "%s%s%s%s" % (
  71. curses.tparm(curses.tigetstr("setaf"), color).decode('ascii')
  72. if color != colors.BLACK else "",
  73. curses.tparm(curses.tigetstr("bold")).decode('ascii')
  74. if bold else "",
  75. string,
  76. curses.tparm(curses.tigetstr("sgr0")).decode('ascii')
  77. )
  78. else:
  79. return string
  80. def move_cursor(direction, amount=1, stream=sys.stdout, erase=False):
  81. """Move the cursor.
  82. If the stream is a TTY, print characters that will move the cursor
  83. in the given direction and optionally erase the line. Else do nothing.
  84. direction (int): the direction as a directions constant, like
  85. directions.UP.
  86. stream (fileobj): a file-like object (that adheres to the API
  87. declared in the `io' package). Defaults to sys.stdout.
  88. erase (bool): True if the line the cursor ends on should be erased.
  89. """
  90. if stream.isatty():
  91. if direction == directions.UP:
  92. print(curses.tparm(curses.tigetstr("cuu"), amount).decode('ascii'),
  93. file=stream, end='')
  94. elif direction == directions.DOWN:
  95. print(curses.tparm(curses.tigetstr("cud"), amount).decode('ascii'),
  96. file=stream, end='')
  97. elif direction == directions.LEFT:
  98. print(curses.tparm(curses.tigetstr("cub"), amount).decode('ascii'),
  99. file=stream, end='')
  100. elif direction == directions.RIGHT:
  101. print(curses.tparm(curses.tigetstr("cuf"), amount).decode('ascii'),
  102. file=stream, end='')
  103. if erase:
  104. print(curses.tparm(curses.tigetstr("el")).decode('ascii'),
  105. file=stream, end='')
  106. stream.flush()