File: term.py

package info (click to toggle)
python-stem 1.8.2-1.1
  • links: PTS, VCS
  • area: main
  • in suites: forky, sid, trixie
  • size: 7,852 kB
  • sloc: python: 33,748; java: 312; makefile: 124; sh: 14
file content (157 lines) | stat: -rw-r--r-- 4,774 bytes parent folder | download | duplicates (3)
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
# Copyright 2011-2019, Damian Johnson and The Tor Project
# See LICENSE for licensing information

"""
Utilities for working with the terminal.

**Module Overview:**

::

  encoding - provides the ANSI escape sequence for a terminal attribute
  format - wrap text with ANSI for the given colors or attributes

.. data:: Color (enum)
.. data:: BgColor (enum)

  Foreground or background terminal colors.

  =========== ===========
  Color       Description
  =========== ===========
  **BLACK**   black color
  **BLUE**    blue color
  **CYAN**    cyan color
  **GREEN**   green color
  **MAGENTA** magenta color
  **RED**     red color
  **WHITE**   white color
  **YELLOW**  yellow color
  =========== ===========

.. data:: Attr (enum)

  Terminal text attributes.

  .. versionchanged:: 1.5.0
     Added the LINES attribute.

  =================== ===========
  Attr                Description
  =================== ===========
  **BOLD**            heavy typeface
  **HIGHLIGHT**       inverted foreground and background
  **UNDERLINE**       underlined text
  **READLINE_ESCAPE** wrap encodings in `RL_PROMPT_START_IGNORE and RL_PROMPT_END_IGNORE sequences <https://stackoverflow.com/questions/9468435/look-how-to-fix-column-calculation-in-python-readline-if-use-color-prompt>`_
  **LINES**           formats lines individually
  =================== ===========
"""

import stem.util.enum
import stem.util.str_tools

TERM_COLORS = ('BLACK', 'RED', 'GREEN', 'YELLOW', 'BLUE', 'MAGENTA', 'CYAN', 'WHITE')

# DISABLE_COLOR_SUPPORT is *not* being vended to Stem users. This is likely to
# go away if I can think of a more graceful method for color toggling.

DISABLE_COLOR_SUPPORT = False

Color = stem.util.enum.Enum(*TERM_COLORS)
BgColor = stem.util.enum.Enum(*['BG_' + color for color in TERM_COLORS])
Attr = stem.util.enum.Enum('BOLD', 'UNDERLINE', 'HIGHLIGHT', 'READLINE_ESCAPE', 'LINES')

# mappings of terminal attribute enums to their ANSI escape encoding
FG_ENCODING = dict([(list(Color)[i], str(30 + i)) for i in range(8)])
BG_ENCODING = dict([(list(BgColor)[i], str(40 + i)) for i in range(8)])
ATTR_ENCODING = {Attr.BOLD: '1', Attr.UNDERLINE: '4', Attr.HIGHLIGHT: '7'}

CSI = '\x1B[%sm'
RESET = CSI % '0'


def encoding(*attrs):
  """
  Provides the ANSI escape sequence for these terminal color or attributes.

  .. versionadded:: 1.5.0

  :param list attr: :data:`~stem.util.terminal.Color`,
    :data:`~stem.util.terminal.BgColor`, or :data:`~stem.util.terminal.Attr` to
    provide an ecoding for

  :returns: **str** of the ANSI escape sequence, **None** no attributes are
    recognized
  """

  term_encodings = []

  for attr in attrs:
    # TODO: Account for an earlier misspelled attribute. This should be dropped
    # in Stem. 2.0.x.

    if attr == 'HILIGHT':
      attr = 'HIGHLIGHT'

    attr = stem.util.str_tools._to_camel_case(attr)
    term_encoding = FG_ENCODING.get(attr, None)
    term_encoding = BG_ENCODING.get(attr, term_encoding)
    term_encoding = ATTR_ENCODING.get(attr, term_encoding)

    if term_encoding:
      term_encodings.append(term_encoding)

  if term_encodings:
    return CSI % ';'.join(term_encodings)


def format(msg, *attr):
  """
  Simple terminal text formatting using `ANSI escape sequences
  <https://en.wikipedia.org/wiki/ANSI_escape_code#CSI_codes>`_.
  The following are some toolkits providing similar capabilities:

  * `django.utils.termcolors <https://github.com/django/django/blob/master/django/utils/termcolors.py>`_
  * `termcolor <https://pypi.org/project/termcolor/>`_
  * `colorama <https://pypi.org/project/colorama/>`_

  .. versionchanged:: 1.6.0
     Normalized return value to be unicode to better support python 2/3
     compatibility.

  :param str msg: string to be formatted
  :param str attr: text attributes, this can be :data:`~stem.util.term.Color`,
    :data:`~stem.util.term.BgColor`, or :data:`~stem.util.term.Attr` enums
    and are case insensitive (so strings like 'red' are fine)

  :returns: **unicode** wrapped with ANSI escape encodings, starting with the given
    attributes and ending with a reset
  """

  msg = stem.util.str_tools._to_unicode(msg)

  if DISABLE_COLOR_SUPPORT:
    return msg

  if Attr.LINES in attr:
    attr = list(attr)
    attr.remove(Attr.LINES)
    lines = [format(line, *attr) for line in msg.split('\n')]
    return '\n'.join(lines)

  # if we have reset sequences in the message then apply our attributes
  # after each of them

  if RESET in msg:
    return ''.join([format(comp, *attr) for comp in msg.split(RESET)])

  prefix, suffix = encoding(*attr), RESET

  if prefix:
    if Attr.READLINE_ESCAPE in attr:
      prefix = '\001%s\002' % prefix
      suffix = '\001%s\002' % suffix

    return prefix + msg + suffix
  else:
    return msg