From 9d796d889b9d8d3973a1bd71ea7b82f5e9cd7faa Mon Sep 17 00:00:00 2001 From: Xuehai Pan Date: Mon, 1 Aug 2022 01:47:09 +0800 Subject: [PATCH] feat(gui): add 256 color support for bar charts Signed-off-by: Xuehai Pan --- README.md | 8 +- nvitop/cli.py | 13 ++- nvitop/gui/library/history.py | 2 + nvitop/gui/library/libcurses.py | 131 +++++++++++++++++++++-------- nvitop/gui/screens/main/device.py | 33 +++++--- nvitop/gui/screens/main/host.py | 53 +++++++----- nvitop/gui/screens/main/process.py | 3 +- 7 files changed, 171 insertions(+), 72 deletions(-) diff --git a/README.md b/README.md index 5456a96..058319b 100644 --- a/README.md +++ b/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. diff --git a/nvitop/cli.py b/nvitop/cli.py index 1541ac1..ce70b3d 100644 --- a/nvitop/cli.py +++ b/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, diff --git a/nvitop/gui/library/history.py b/nvitop/gui/library/history.py index 42fb2ea..060c94f 100644 --- a/nvitop/gui/library/history.py +++ b/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: diff --git a/nvitop/gui/library/libcurses.py b/nvitop/gui/library/libcurses.py index 7f9d2d0..a3644fa 100644 --- a/nvitop/gui/library/libcurses.py +++ b/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: diff --git a/nvitop/gui/screens/main/device.py b/nvitop/gui/screens/main/device.py index 241f851..37a4c9d 100644 --- a/nvitop/gui/screens/main/device.py +++ b/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 diff --git a/nvitop/gui/screens/main/host.py b/nvitop/gui/screens/main/host.py index 45ad781..4bd2e49 100644 --- a/nvitop/gui/screens/main/host.py +++ b/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() diff --git a/nvitop/gui/screens/main/process.py b/nvitop/gui/screens/main/process.py index 6c108c8..de61079 100644 --- a/nvitop/gui/screens/main/process.py +++ b/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')