add custom color

wookayin · Jul 18, 2022 · 4780f9b · 4780f9b
1 parent 211346a
commit 4780f9b
Show file tree

Hide file tree

Showing 3 changed files with 48 additions and 13 deletions.
diff --git a/gpustat/cli.py b/gpustat/cli.py
@@ -9,6 +9,7 @@
 from blessed import Terminal
 
 from gpustat import __version__
+from gpustat.custom_color_action import CustomColorAction
 from .core import GPUStatCollection
 
 
@@ -78,6 +79,8 @@ def main(*argv):
                               help='Force to output with colors')
     parser_color.add_argument('--no-color', action='store_true',
                               help='Suppress colored output')
+    parser_color.add_argument('--custom-colors', action=CustomColorAction,
+                              help='Specify output colors')
     parser.add_argument('-a', '--show-all', action='store_true',
                         help='Display all gpu properties above')
 

diff --git a/gpustat/core.py b/gpustat/core.py
@@ -36,6 +36,7 @@
 
 
 class GPUStat(object):
+    CUSTOM_COLORS = {}
 
     def __init__(self, entry):
         if not isinstance(entry, dict):
@@ -186,13 +187,17 @@ def print_to(self, fp,
                  show_power=None,
                  gpuname_width=16,
                  term=None,
+                 custom_colors=None,
                  ):
         if term is None:
             term = Terminal(stream=sys.stdout)
 
         # color settings
         colors = {}
 
+        if custom_colors is None:
+            custom_colors = {}
+
         def _conditional(cond_fn, true_value, false_value,
                          error_value=term.bold_black):
             try:
@@ -202,19 +207,30 @@ def _conditional(cond_fn, true_value, false_value,
 
         _ENC_THRESHOLD = 50
 
-        colors['C0'] = term.normal
-        colors['C1'] = term.cyan
-        colors['CBold'] = term.bold
-        colors['CName'] = term.blue
+        def _set_color(name: str, default_color: str):
+            color = custom_colors.get(name, '')
+            if name in GPUStat.CUSTOM_COLORS:
+                return GPUStat.CUSTOM_COLORS[name]
+            elif color and getattr(term, color):
+                GPUStat.CUSTOM_COLORS[name] = getattr(term, color)
+                return GPUStat.CUSTOM_COLORS[name]
+            else:
+                GPUStat.CUSTOM_COLORS[name] = getattr(term, default_color)
+                return GPUStat.CUSTOM_COLORS[name]
+
+        colors['C0'] = _set_color('C0', 'normal')
+        colors['C1'] = _set_color('C1', 'cyan')
+        colors['CBold'] = _set_color('CBold', 'bold')
+        colors['CName'] = _set_color('CName', 'blue')
         colors['CTemp'] = _conditional(lambda: self.temperature < 50,
                                        term.red, term.bold_red)
         colors['FSpeed'] = _conditional(lambda: self.fan_speed < 30,
                                         term.cyan, term.bold_cyan)
-        colors['CMemU'] = term.bold_yellow
-        colors['CMemT'] = term.yellow
-        colors['CMemP'] = term.yellow
-        colors['CCPUMemU'] = term.yellow
-        colors['CUser'] = term.bold_black   # gray
+        colors['CMemU'] = _set_color('CMemU', 'bold_yellow')
+        colors['CMemT'] = _set_color('CMemT', 'yellow')
+        colors['CMemP'] = _set_color('CMemP', 'yellow')
+        colors['CCPUMemU'] = _set_color('CCPUMemU', 'yellow')
+        colors['CUser'] = _set_color('CUser', 'bold_black')
         colors['CUtil'] = _conditional(lambda: self.utilization < 30,
                                        term.green, term.bold_green)
         colors['CUtilEnc'] = _conditional(
@@ -223,13 +239,13 @@ def _conditional(cond_fn, true_value, false_value,
         colors['CUtilDec'] = _conditional(
             lambda: self.utilization_dec < _ENC_THRESHOLD,
             term.green, term.bold_green)
-        colors['CCPUUtil'] = term.green
+        colors['CCPUUtil'] = _set_color('CCPUUtil', 'green')
         colors['CPowU'] = _conditional(
             lambda: (self.power_limit is not None and
                      float(self.power_draw) / self.power_limit < 0.4),
             term.magenta, term.bold_magenta
         )
-        colors['CPowL'] = term.magenta
+        colors['CPowL'] = _set_color('CPowL', 'magenta')
         colors['CCmd'] = term.color(24)   # a bit dark
 
         if not with_colors:
@@ -574,7 +590,7 @@ def print_formatted(self, fp=sys.stdout, force_color=False, no_color=False,
                         show_pid=False, show_fan_speed=None,
                         show_codec="", show_power=None,
                         gpuname_width=16, show_header=True,
-                        eol_char=os.linesep,
+                        eol_char=os.linesep, custom_colors=None
                         ):
         # ANSI color configuration
         if force_color and no_color:
@@ -631,7 +647,8 @@ def print_formatted(self, fp=sys.stdout, force_color=False, no_color=False,
                        show_codec=show_codec,
                        show_power=show_power,
                        gpuname_width=gpuname_width,
-                       term=t_color)
+                       term=t_color,
+                       custom_colors=custom_colors)
             fp.write(eol_char)
 
         fp.flush()

diff --git a/gpustat/custom_color_action.py b/gpustat/custom_color_action.py
@@ -0,0 +1,15 @@
+import argparse
+
+
+class CustomColorAction(argparse.Action):
+    def __call__(self, _parser, namespace, values, _option_strings=None):
+        param_dict = getattr(namespace, self.dest)
+        if param_dict is None:
+            param_dict = {}
+
+        try:
+            k, v = values.split("=")
+            param_dict[k] = v
+        except ValueError:
+            pass
+        setattr(namespace, self.dest, param_dict)