From 6acbd4be827a59d7e0a80ccd54b880e5837099a8 Mon Sep 17 00:00:00 2001
From: Oleksandr Kuvshynov <661042+okuvshynov@users.noreply.github.com>
Date: Tue, 10 Sep 2024 11:49:35 -0400
Subject: [PATCH] cubestat: fix colors for non-256 color terminals

---
 cubestat/colors.py                   | 46 ++++++++++++++++++++++------
 cubestat/screen.py                   |  6 ++--
 cubestat/tests/test_prepare_cells.py |  6 ++--
 3 files changed, 42 insertions(+), 16 deletions(-)

diff --git a/cubestat/colors.py b/cubestat/colors.py
index 7dfa4ee..33cfcb2 100644
--- a/cubestat/colors.py
+++ b/cubestat/colors.py
@@ -1,4 +1,5 @@
 import curses
+import logging
 
 from cubestat.common import DisplayMode
 
@@ -30,18 +31,26 @@ class ColorTheme(DisplayMode):
 }
 
 
-def prepare_cells():
+def mono_cells():
     chrs = [' ', '▁', '▂', '▃', '▄', '▅', '▆', '▇', '█']
-    cells = {}
-    colorpair = 1
-    for name, colors in colors_ansi256.items():
-        cells[name] = []
-        for i, (fg, bg) in enumerate(zip(colors[1:], colors[:-1])):
+    return [(c, 0) for c in chrs]
+
+
+def cells_for_colorscheme(colors, colorpair):
+    chrs = [' ', '▁', '▂', '▃', '▄', '▅', '▆', '▇', '█']
+    res = []
+    for i, (fg, bg) in enumerate(zip(colors[1:], colors[:-1])):
+        try:
             curses.init_pair(colorpair, fg, bg)
-            j = 0 if i == 0 else 1
-            cells[name].extend((chr, colorpair) for chr in chrs[j:])
-            colorpair += 1
-    return cells
+        except curses.error:
+            logging.error('curses.init_pair returned error.')
+            logging.error('Consider setting up 256 colors terminal:')
+            logging.error('  export TERM=xterm-256color')
+            return None, colorpair
+        j = 0 if i == 0 else 1
+        res.extend((chr, colorpair) for chr in chrs[j:])
+        colorpair += 1
+    return res, colorpair
 
 
 def get_theme(metric, color_mode):
@@ -51,3 +60,20 @@ def get_theme(metric, color_mode):
         ColorTheme.col:  light_colormap.get(metric, 'green')
     }
     return res[color_mode]
+
+
+class Colorschemes:
+    def __init__(self):
+        self.schemes = {}
+        colorpair = 1
+        for name, colors in colors_ansi256.items():
+            cells, colorpair = cells_for_colorscheme(colors, colorpair)
+            if cells is not None:
+                self.schemes[name] = cells
+        self.schemes['mono'] = mono_cells()
+
+    def get_cells(self, name):
+        if name in self.schemes:
+            return self.schemes[name]
+        logging.error(f'requested colorscheme {name}.')
+        return self.schemes['mono']
diff --git a/cubestat/screen.py b/cubestat/screen.py
index debfe7d..b8c8ebf 100644
--- a/cubestat/screen.py
+++ b/cubestat/screen.py
@@ -1,7 +1,7 @@
 import curses
 from math import floor
 
-from cubestat.colors import prepare_cells
+from cubestat.colors import Colorschemes
 
 
 class Screen:
@@ -13,7 +13,7 @@ def __init__(self, stdscr):
         curses.start_color()
         curses.use_default_colors()
         self.spacing = ' '
-        self.cells = prepare_cells()
+        self.colors = Colorschemes()
 
     def write_string(self, row, col, s, color=0):
         if col + len(s) > self.cols:
@@ -54,7 +54,7 @@ def render_ruler(self, indent, title, base_ruler, items, row):
         self.write_string(row + 1, self.cols - len(bottomright_border), bottomright_border)
 
     def render_chart(self, theme, max_value, data, row):
-        cells = self.cells[theme]
+        cells = self.colors.get_cells(theme)
         scaler = len(cells) / max_value
         col_start = self.cols - (len(data) + len(self.spacing)) - 1
 
diff --git a/cubestat/tests/test_prepare_cells.py b/cubestat/tests/test_prepare_cells.py
index 9de4382..530c002 100644
--- a/cubestat/tests/test_prepare_cells.py
+++ b/cubestat/tests/test_prepare_cells.py
@@ -1,4 +1,4 @@
-from cubestat.colors import prepare_cells
+from cubestat.colors import Colorschemes
 
 from unittest.mock import patch
 import unittest
@@ -7,8 +7,8 @@
 class TestPrepareCells(unittest.TestCase):
     @patch('curses.init_pair')
     def test_prepare_cells(self, mock_init_pair):
-        cells = prepare_cells()
-        self.assertEqual(len(cells['green']), 25)
+        colors = Colorschemes()
+        self.assertEqual(len(colors.schemes['green']), 25)
         self.assertEqual(mock_init_pair.call_count, 18)