feat(gui): add 256 color support for bar charts

Signed-off-by: Xuehai Pan <XuehaiPan@pku.edu.cn>

README.md

@@ -69,7 +69,7 @@ If this repo is useful to you, please star ⭐️ it to let more people know
 
 - **Informative and fancy output**: show more information than `nvidia-smi` with colorized fancy box drawing.
 - **Monitor mode**: can run as a resource monitor, rather than print the results only once. (vs. [nvidia-htop](https://github.com/peci1/nvidia-htop), limited support with command `watch -c`)
-  - bar plots and history graphs
+  - bar charts and history graphs
   - process sorting
   - process filtering
   - send signals to processes with a keystroke
@@ -278,7 +278,7 @@ Type `nvitop --help` for more command options:
 
 ```text
 usage: nvitop [--help] [--version] [--once] [--monitor [{auto,full,compact}]]
-              [--interval SEC] [--ascii] [--force-color] [--light]
+              [--interval SEC] [--ascii] [--colorful] [--force-color] [--light]
               [--gpu-util-thresh th1 th2] [--mem-util-thresh th1 th2]
               [--only idx [idx ...]] [--only-visible] [--compute] [--graphics]
               [--user [USERNAME ...]] [--pid PID [PID ...]]
@@ -298,6 +298,10 @@ optional arguments:
                         Use ASCII characters only, which is useful for terminals without Unicode support.
 
 coloring:
+  --colorful            Use gradient colors to get spectrum-like bar charts. This option is only available
+                        when the terminal supports 256 colors. You may need to set environment variable
+                        `TERM="xterm-256color"`. Note that the terminal multiplexer, such as `tmux`, may
+                        override the `TREM` variable.
   --force-color         Force colorize even when `stdout` is not a TTY terminal.
   --light               Tweak visual results for light theme terminals in monitor mode.
                         Set variable `NVITOP_MONITOR_THEME="light"` on light terminals for convenience.

nvitop/cli.py

@@ -87,6 +87,17 @@ def parse_arguments():  # pylint: disable=too-many-branches,too-many-statements
     )
 
     coloring = parser.add_argument_group('coloring')
+    coloring.add_argument(
+        '--colorful',
+        dest='colorful',
+        action='store_true',
+        help=(
+            'Use gradient colors to get spectrum-like bar charts. This option is only available\n'
+            'when the terminal supports 256 colors. You may need to set environment variable\n'
+            '`TERM="xterm-256color"`. Note that the terminal multiplexer, such as `tmux`, may\n'
+            'override the `TREM` variable.'
+        ),
+    )
     coloring.add_argument(
         '--force-color',
         dest='force_color',
@@ -298,7 +309,7 @@ def main():  # pylint: disable=too-many-branches,too-many-statements,too-many-lo
     top = None
     if hasattr(args, 'monitor') and len(devices) > 0:
         try:
-            with libcurses(light_theme=args.light) as win:
+            with libcurses(colorful=args.colorful, light_theme=args.light) as win:
                 top = Top(
                     devices,
                     filters,

nvitop/gui/library/history.py

@@ -156,6 +156,8 @@ class HistoryGraph:  # pylint: disable=too-many-instance-attributes
         except ValueError:
             return NA
 
+    __str__ = last_value_string
+
     def max_value_string(self):
         max_value = self.max_value
         if max_value >= self.baseline:

nvitop/gui/library/libcurses.py

@@ -3,6 +3,7 @@
 
 # pylint: disable=missing-module-docstring,missing-class-docstring,missing-function-docstring
 
+import colorsys
 import contextlib
 import curses
 import locale
@@ -14,6 +15,27 @@ LIGHT_THEME = False
 DEFAULT_FOREGROUND = curses.COLOR_WHITE
 DEFAULT_BACKGROUND = curses.COLOR_BLACK
 COLOR_PAIRS = {None: 0}
+TRUE_COLORS = dict(
+    [
+        ('black', 0),
+        ('red', 1),
+        ('green', 2),
+        ('yellow', 3),
+        ('blue', 4),
+        ('magenta', 5),
+        ('cyan', 6),
+        ('white', 7),
+        ('bright black', 8),
+        ('bright red', 9),
+        ('bright green', 10),
+        ('bright yellow', 11),
+        ('bright blue', 12),
+        ('bright magenta', 13),
+        ('bright cyan', 14),
+        ('bright white', 15),
+    ]
+    + [('preserved {:02d}'.format(i), i) for i in range(16, 48)]
+)
 
 
 def _init_color_theme(light_theme=False):
@@ -30,6 +52,24 @@ def _init_color_theme(light_theme=False):
         DEFAULT_BACKGROUND = curses.COLOR_BLACK
 
 
+def _colormap(x, levels=200):
+    # pylint: disable=invalid-name
+    h = 0.5 * (1.0 - x) - 0.15
+    h = (round(h * levels) / levels) % 1.0
+    r, g, b = colorsys.hsv_to_rgb(h, 0.7, 0.8)
+    return (round(1000.0 * r), round(1000.0 * g), round(1000.0 * b))
+
+
+def _get_true_color(rgb):
+    if rgb not in TRUE_COLORS:
+        try:
+            curses.init_color(len(TRUE_COLORS), *rgb)
+        except curses.error:
+            return -1
+        TRUE_COLORS[rgb] = len(TRUE_COLORS)
+    return TRUE_COLORS[rgb]
+
+
 def _get_color(fg, bg):
     """Returns the curses color pair for the given fg/bg combination."""
 
@@ -37,8 +77,16 @@ def _get_color(fg, bg):
 
     if isinstance(fg, str):
         fg = getattr(curses, 'COLOR_{}'.format(fg.upper()), -1)
+    elif isinstance(fg, tuple):
+        fg = _get_true_color(fg)
+    elif isinstance(fg, float):
+        fg = _get_true_color(_colormap(fg))
     if isinstance(bg, str):
         bg = getattr(curses, 'COLOR_{}'.format(bg.upper()), -1)
+    elif isinstance(bg, tuple):
+        bg = _get_true_color(bg)
+    elif isinstance(bg, float):
+        bg = _get_true_color(_colormap(bg))
 
     key = (fg, bg)
     if key not in COLOR_PAIRS:
@@ -64,23 +112,6 @@ def _get_color(fg, bg):
     return COLOR_PAIRS[key]
 
 
-def _get_color_attr(fg=-1, bg=-1, attr=0):
-    """Returns the curses attribute for the given fg/bg/attr combination."""
-    if isinstance(attr, str):
-        attr_strings = map(str.strip, attr.split('|'))
-        attr = 0
-        for s in attr_strings:
-            attr |= getattr(curses, 'A_{}'.format(s.upper()), 0)
-
-    if LIGHT_THEME:  # tweak for light themes
-        if attr & curses.A_REVERSE != 0 and bg == -1 and fg not in (DEFAULT_FOREGROUND, -1):
-            bg = DEFAULT_FOREGROUND
-
-    if fg == -1 and bg == -1:
-        return attr
-    return curses.color_pair(_get_color(fg, bg)) | attr
-
-
 def setlocale_utf8():
     for code in ('C.UTF-8', 'en_US.UTF-8', '', 'C'):
         try:
@@ -95,7 +126,7 @@ def setlocale_utf8():
 
 
 @contextlib.contextmanager
-def libcurses(light_theme=False):
+def libcurses(colorful=False, light_theme=False):
     os.environ.setdefault('ESCDELAY', '25')
     setlocale_utf8()
 
@@ -119,6 +150,12 @@ def libcurses(light_theme=False):
     except curses.error:
         pass
 
+    if colorful:
+        try:
+            CursesShortcuts.TERM_256COLOR = curses.COLORS >= 256
+        except AttributeError:
+            pass
+
     # Push a Ctrl+C (ascii value 3) to the curses getch stack
     def interrupt_handler(signalnum, frame):  # pylint: disable=unused-argument
         curses.ungetch(3)
@@ -146,9 +183,10 @@ class CursesShortcuts:
         '═' + '─╴' + '╒╤╕╪╘╧╛┌┬┐┼└┴┘' + '│╞╡├┤▏▎▍▌▋▊▉█░' + '▲▼' + '␤',
         '=' + '--' + '++++++++++++++' + '||||||||||||||' + '^v' + '?',
     )
+    TERM_256COLOR = False
 
     def __init__(self):
-        self.win = None
+        self.win = None  # type: curses._CursesWindow
         self.ascii = False
 
     def addstr(self, *args, **kwargs):
@@ -189,24 +227,51 @@ class CursesShortcuts:
 
     def color(self, fg=-1, bg=-1, attr=0):
         """Change the colors from now on."""
-        self.set_fg_bg_attr(fg, bg, attr)
 
-    def color_at(self, y, x, width, *args, **kwargs):
-        """Change the colors at the specified position"""
-        try:
-            self.win.chgat(y, x, width, _get_color_attr(*args, **kwargs))
-        except curses.error:
-            pass
-
-    def set_fg_bg_attr(self, fg=-1, bg=-1, attr=0):
-        try:
-            self.win.attrset(_get_color_attr(fg, bg, attr))
-        except curses.error:
-            pass
+        return self.set_fg_bg_attr(fg, bg, attr)
 
     def color_reset(self):
         """Change the colors to the default colors"""
-        self.color()
+
+        return self.color()
+
+    def color_at(self, y, x, width, *args, **kwargs):
+        """Change the colors at the specified position"""
+
+        try:
+            self.win.chgat(y, x, width, self.get_fg_bg_attr(*args, **kwargs))
+        except curses.error:
+            pass
+
+    @staticmethod
+    def get_fg_bg_attr(fg=-1, bg=-1, attr=0):
+        """Returns the curses attribute for the given fg/bg/attr combination."""
+
+        if fg == -1 and bg == -1 and attr == 0:
+            return 0
+
+        if isinstance(attr, str):
+            attr_strings = map(str.strip, attr.split('|'))
+            attr = 0
+            for s in attr_strings:
+                attr |= getattr(curses, 'A_{}'.format(s.upper()), 0)
+
+        if LIGHT_THEME:  # tweak for light themes
+            if attr & curses.A_REVERSE != 0 and bg == -1 and fg not in (DEFAULT_FOREGROUND, -1):
+                bg = DEFAULT_FOREGROUND
+
+        if fg == -1 and bg == -1:
+            return attr
+        return curses.color_pair(_get_color(fg, bg)) | attr
+
+    def set_fg_bg_attr(self, fg=-1, bg=-1, attr=0):
+        try:
+            attr = self.get_fg_bg_attr(fg, bg, attr)
+            self.win.attrset(attr)
+        except curses.error:
+            return 0
+        else:
+            return attr
 
     def update_size(self, termsize=None):
         if termsize is None:

nvitop/gui/screens/main/device.py

@@ -254,7 +254,8 @@ class DevicePanel(Displayable):  # pylint: disable=too-many-instance-attributes
 
         super().poke()
 
-    def draw(self):  # pylint: disable=too-many-locals,too-many-branches,too-many-statements
+    # pylint: disable-next=too-many-locals,too-many-branches,too-many-statements
+    def draw(self):
         self.color_reset()
 
         if self.need_redraw:
@@ -280,7 +281,7 @@ class DevicePanel(Displayable):  # pylint: disable=too-many-instance-attributes
 
         y_start = self.y + len(formats) + 5
         prev_device_index = self.snapshots[0].tuple_index
-        for index, device in enumerate(self.snapshots):
+        for index, device in enumerate(self.snapshots):  # pylint: disable=too-many-nested-blocks
             if (
                 len(prev_device_index) != len(device.tuple_index)
                 or prev_device_index[0] != device.tuple_index[0]
@@ -306,7 +307,7 @@ class DevicePanel(Displayable):  # pylint: disable=too-many-instance-attributes
             if draw_bars:
                 matrix = [
                     (
-                        0,
+                        self.x + 80,
                         y_start,
                         remaining_width - 3,
                         'MEM',
@@ -314,7 +315,7 @@ class DevicePanel(Displayable):  # pylint: disable=too-many-instance-attributes
                         device.memory_display_color,
                     ),
                     (
-                        0,
+                        self.x + 80,
                         y_start + 1,
                         remaining_width - 3,
                         'UTL',
@@ -328,7 +329,7 @@ class DevicePanel(Displayable):  # pylint: disable=too-many-instance-attributes
                         right_width = (remaining_width - 6) // 2 + 1
                         matrix = [
                             (
-                                0,
+                                self.x + 80,
                                 y_start,
                                 left_width,
                                 'MEM',
@@ -336,7 +337,7 @@ class DevicePanel(Displayable):  # pylint: disable=too-many-instance-attributes
                                 device.memory_display_color,
                             ),
                             (
-                                left_width + 3,
+                                self.x + 80 + left_width + 3,
                                 y_start,
                                 right_width,
                                 'UTL',
@@ -358,11 +359,21 @@ class DevicePanel(Displayable):  # pylint: disable=too-many-instance-attributes
                 elif device.is_mig_device:
                     matrix.pop()
                 for x_offset, y, width, prefix, utilization, color in matrix:
-                    bar = make_bar(  # pylint: disable=disallowed-name
-                        prefix, utilization, width=width
-                    )
-                    self.addstr(y, self.x + 80 + x_offset, bar)
-                    self.color_at(y, self.x + 80 + x_offset, width=width, fg=color, attr=attr)
+                    # pylint: disable-next=disallowed-name
+                    bar = make_bar(prefix, utilization, width=width)
+                    self.addstr(y, x_offset, bar)
+                    if self.TERM_256COLOR:
+                        parts = bar.rstrip().split(' ')
+                        prefix_len = len(parts[0])
+                        bar_len = len(parts[1])
+                        full_bar_len = width - prefix_len - 5
+                        self.color_at(y, x_offset, width=width, fg=float(bar_len / full_bar_len))
+                        for i, x in enumerate(
+                            range(x_offset + prefix_len + 1, x_offset + prefix_len + 1 + bar_len)
+                        ):
+                            self.color_at(y, x, width=1, fg=float(i / full_bar_len))
+                    else:
+                        self.color_at(y, x_offset, width=width, fg=color, attr=attr)
 
             y_start += len(fmts)
             prev_device_index = device.tuple_index

nvitop/gui/screens/main/host.py

@@ -296,26 +296,17 @@ class HostPanel(Displayable):  # pylint: disable=too-many-instance-attributes
                         attr='dim',
                     )
 
+        self.color(fg='cyan')
         for y, line in enumerate(host.cpu_percent.history.graph, start=self.y):
             self.addstr(y, self.x + 1, line)
-            self.color_at(y, self.x + 1, width=77, fg='cyan')
-        self.addstr(self.y, self.x + 1, ' {} '.format(load_average))
-        self.addstr(
-            self.y + 1, self.x + 1, ' {} '.format(host.cpu_percent.history.last_value_string())
-        )
 
+        self.color(fg='magenta')
         for y, line in enumerate(host.memory_percent.history.graph, start=self.y + 6):
             self.addstr(y, self.x + 1, line)
-            self.color_at(y, self.x + 1, width=77, fg='magenta')
-        self.addstr(
-            self.y + 9, self.x + 1, ' {} '.format(host.memory_percent.history.last_value_string())
-        )
+
+        self.color(fg='blue')
         for y, line in enumerate(host.swap_percent.history.graph, start=self.y + 10):
             self.addstr(y, self.x + 1, line)
-            self.color_at(y, self.x + 1, width=77, fg='blue')
-        self.addstr(
-            self.y + 10, self.x + 1, ' {} '.format(host.swap_percent.history.last_value_string())
-        )
 
         if self.width >= 100:
             if self.device_count > 1 and self.parent.selected.is_set():
@@ -326,20 +317,36 @@ class HostPanel(Displayable):  # pylint: disable=too-many-instance-attributes
                 memory_percent = self.average_memory_percent
                 gpu_utilization = self.average_gpu_utilization
 
-            memory_display_color = Device.color_of(memory_percent.last_value, type='memory')
-            gpu_display_color = Device.color_of(gpu_utilization.last_value, type='gpu')
+            if self.TERM_256COLOR:
+                for i, (y, line) in enumerate(enumerate(memory_percent.graph, start=self.y)):
+                    self.addstr(y, self.x + 79, line, self.get_fg_bg_attr(fg=1.0 - i / 4.0))
 
-            for y, line in enumerate(memory_percent.graph, start=self.y):
-                self.addstr(y, self.x + 79, line)
-                self.color_at(y, self.x + 79, width=remaining_width - 1, fg=memory_display_color)
-            self.addstr(self.y, self.x + 79, ' {} '.format(memory_percent.last_value_string()))
+                for i, (y, line) in enumerate(enumerate(gpu_utilization.graph, start=self.y + 6)):
+                    self.addstr(y, self.x + 79, line, self.get_fg_bg_attr(fg=i / 4.0))
+            else:
+                self.color(fg=Device.color_of(memory_percent.last_value, type='memory'))
+                for y, line in enumerate(memory_percent.graph, start=self.y):
+                    self.addstr(y, self.x + 79, line)
 
-            for y, line in enumerate(gpu_utilization.graph, start=self.y + 6):
-                self.addstr(y, self.x + 79, line)
-                self.color_at(y, self.x + 79, width=remaining_width - 1, fg=gpu_display_color)
+                self.color(fg=Device.color_of(gpu_utilization.last_value, type='gpu'))
+                for y, line in enumerate(gpu_utilization.graph, start=self.y + 6):
+                    self.addstr(y, self.x + 79, line)
+
+            self.color_reset()
+            self.addstr(self.y, self.x + 1, ' {} '.format(load_average))
+            self.addstr(self.y + 1, self.x + 1, ' {} '.format(host.cpu_percent.history))
             self.addstr(
-                self.y + 10, self.x + 79, ' {} '.format(gpu_utilization.last_value_string())
+                self.y + 9,
+                self.x + 1,
+                ' {} '.format(host.memory_percent.history),
             )
+            self.addstr(
+                self.y + 10,
+                self.x + 1,
+                ' {} '.format(host.swap_percent.history),
+            )
+            self.addstr(self.y, self.x + 79, ' {} '.format(memory_percent))
+            self.addstr(self.y + 10, self.x + 79, ' {} '.format(gpu_utilization))
 
     def destroy(self):
         super().destroy()

nvitop/gui/screens/main/process.py

@@ -483,8 +483,7 @@ class ProcessPanel(Displayable):  # pylint: disable=too-many-instance-attributes
                     )
                 else:
                     if self.selected.is_same_on_host(process):
-                        self.addstr(y, self.x + 1, '=')
-                        self.color_at(y, self.x + 1, width=1, attr='bold | blink')
+                        self.addstr(y, self.x + 1, '=', self.get_fg_bg_attr(attr='bold | blink'))
                     self.color_at(y, self.x + 2, width=3, fg=color)
                     if str(process.username) != USERNAME and not SUPERUSER:
                         self.color_at(y, self.x + 5, width=self.width - 6, attr='dim')