You can not select more than 25 topics Topics must start with a letter or number, can include dashes ('-') and can be up to 35 characters long.

222 lines
6.5 KiB

4 years ago
  1. # coding: utf8
  2. from __future__ import unicode_literals, print_function
  3. import os
  4. import sys
  5. import textwrap
  6. import difflib
  7. import itertools
  8. STDOUT_ENCODING = sys.stdout.encoding if hasattr(sys.stdout, "encoding") else None
  9. ENCODING = STDOUT_ENCODING or "ascii"
  10. NO_UTF8 = ENCODING.lower() not in ("utf8", "utf-8")
  11. # Environment variables
  12. ENV_ANSI_DISABLED = "ANSI_COLORS_DISABLED" # no colors
  13. class MESSAGES(object):
  14. GOOD = "good"
  15. FAIL = "fail"
  16. WARN = "warn"
  17. INFO = "info"
  18. COLORS = {
  19. MESSAGES.GOOD: 2,
  20. MESSAGES.FAIL: 1,
  21. MESSAGES.WARN: 3,
  22. MESSAGES.INFO: 4,
  23. "red": 1,
  24. "green": 2,
  25. "yellow": 3,
  26. "blue": 4,
  27. "pink": 5,
  28. "cyan": 6,
  29. "white": 7,
  30. "grey": 8,
  31. "black": 16,
  32. }
  33. ICONS = {
  34. MESSAGES.GOOD: "\u2714" if not NO_UTF8 else "[+]",
  35. MESSAGES.FAIL: "\u2718" if not NO_UTF8 else "[x]",
  36. MESSAGES.WARN: "\u26a0" if not NO_UTF8 else "[!]",
  37. MESSAGES.INFO: "\u2139" if not NO_UTF8 else "[i]",
  38. }
  39. # Python 2 compatibility
  40. IS_PYTHON_2 = sys.version_info[0] == 2
  41. if IS_PYTHON_2:
  42. basestring_ = basestring # noqa: F821
  43. input_ = raw_input # noqa: F821
  44. zip_longest = itertools.izip_longest # noqa: F821
  45. else:
  46. basestring_ = str
  47. input_ = input
  48. zip_longest = itertools.zip_longest
  49. def color(text, fg=None, bg=None, bold=False, underline=False):
  50. """Color text by applying ANSI escape sequence.
  51. text (unicode): The text to be formatted.
  52. fg (unicode / int): Foreground color. String name or 0 - 256 (see COLORS).
  53. bg (unicode / int): Background color. String name or 0 - 256 (see COLORS).
  54. bold (bool): Format text in bold.
  55. underline (bool): Underline text.
  56. RETURNS (unicode): The formatted text.
  57. """
  58. fg = COLORS.get(fg, fg)
  59. bg = COLORS.get(bg, bg)
  60. if not any([fg, bg, bold]):
  61. return text
  62. styles = []
  63. if bold:
  64. styles.append("1")
  65. if underline:
  66. styles.append("4")
  67. if fg:
  68. styles.append("38;5;{}".format(fg))
  69. if bg:
  70. styles.append("48;5;{}".format(bg))
  71. return "\x1b[{}m{}\x1b[0m".format(";".join(styles), text)
  72. def wrap(text, wrap_max=80, indent=4):
  73. """Wrap text at given width using textwrap module.
  74. text (unicode): The text to wrap.
  75. wrap_max (int): Maximum line width, including indentation. Defaults to 80.
  76. indent (int): Number of spaces used for indentation. Defaults to 4.
  77. RETURNS (unicode): The wrapped text with line breaks.
  78. """
  79. indent = indent * " "
  80. wrap_width = wrap_max - len(indent)
  81. text = to_string(text)
  82. return textwrap.fill(
  83. text,
  84. width=wrap_width,
  85. initial_indent=indent,
  86. subsequent_indent=indent,
  87. break_long_words=False,
  88. break_on_hyphens=False,
  89. )
  90. def format_repr(obj, max_len=50, ellipsis="..."):
  91. """Wrapper around `repr()` to print shortened and formatted string version.
  92. obj: The object to represent.
  93. max_len (int): Maximum string length. Longer strings will be cut in the
  94. middle so only the beginning and end is displayed, separated by ellipsis.
  95. ellipsis (unicode): Ellipsis character(s), e.g. "...".
  96. RETURNS (unicode): The formatted representation.
  97. """
  98. string = repr(obj)
  99. if len(string) >= max_len:
  100. half = int(max_len / 2)
  101. return "{} {} {}".format(string[:half], ellipsis, string[-half:])
  102. else:
  103. return string
  104. def diff_strings(a, b, fg="black", bg=("green", "red")):
  105. """Compare two strings and return a colored diff with red/green background
  106. for deletion and insertions.
  107. a (unicode): The first string to diff.
  108. b (unicode): The second string to diff.
  109. fg (unicode / int): Foreground color. String name or 0 - 256 (see COLORS).
  110. bg (tuple): Background colors as (insert, delete) tuple of string name or
  111. 0 - 256 (see COLORS).
  112. RETURNS (unicode): The formatted diff.
  113. """
  114. output = []
  115. matcher = difflib.SequenceMatcher(None, a, b)
  116. for opcode, a0, a1, b0, b1 in matcher.get_opcodes():
  117. if opcode == "equal":
  118. output.append(a[a0:a1])
  119. elif opcode == "insert":
  120. output.append(color(b[b0:b1], fg=fg, bg=bg[0]))
  121. elif opcode == "delete":
  122. output.append(color(a[a0:a1], fg=fg, bg=bg[1]))
  123. return "".join(output)
  124. def get_raw_input(description, default=False, indent=4):
  125. """Get user input from the command line via raw_input / input.
  126. description (unicode): Text to display before prompt.
  127. default (unicode or False/None): Default value to display with prompt.
  128. indent (int): Indentation in spaces.
  129. RETURNS (unicode): User input.
  130. """
  131. additional = " (default: {})".format(default) if default else ""
  132. prompt = wrap("{}{}: ".format(description, additional), indent=indent)
  133. user_input = input_(prompt)
  134. return user_input
  135. def locale_escape(string, errors="replace"):
  136. """Mangle non-supported characters, for savages with ASCII terminals.
  137. string (unicode): The string to escape.
  138. errors (unicode): The str.encode errors setting. Defaults to `"replace"`.
  139. RETURNS (unicode): The escaped string.
  140. """
  141. string = to_string(string)
  142. string = string.encode(ENCODING, errors).decode("utf8")
  143. return string
  144. def can_render(string):
  145. """Check if terminal can render unicode characters, e.g. special loading
  146. icons. Can be used to display fallbacks for ASCII terminals.
  147. string (unicode): The string to render.
  148. RETURNS (bool): Whether the terminal can render the text.
  149. """
  150. try:
  151. string.encode(ENCODING)
  152. return True
  153. except UnicodeEncodeError:
  154. return False
  155. def supports_ansi():
  156. """Returns True if the running system's terminal supports ANSI escape
  157. sequences for color, formatting etc. and False otherwise. Inspired by
  158. Django's solution – hacky, but an okay approximation.
  159. RETURNS (bool): Whether the terminal supports ANSI colors.
  160. """
  161. if os.getenv(ENV_ANSI_DISABLED):
  162. return False
  163. # See: https://stackoverflow.com/q/7445658/6400719
  164. supported_platform = sys.platform != "Pocket PC" and (
  165. sys.platform != "win32" or "ANSICON" in os.environ
  166. )
  167. if not supported_platform:
  168. return False
  169. return True
  170. def to_string(text):
  171. """Minimal compat helper to make sure text is unicode. Mostly used to
  172. convert Paths and other Python objects.
  173. text: The text/object to be converted.
  174. RETURNS (unicode): The converted string.
  175. """
  176. if not isinstance(text, basestring_):
  177. if IS_PYTHON_2:
  178. text = str(text).decode("utf8")
  179. else:
  180. text = str(text)
  181. return text