From 3f208b4c1c999ab6a0650b170eecd49456d84aa5 Mon Sep 17 00:00:00 2001
From: The MMGen Project <mmgen@tuta.io>
Date: Thu, 27 Jan 2022 11:08:05 +0000
Subject: [PATCH] improve color initialization

---
 mmgen/color.py    | 51 ++++++++++++++++++++---------------------------
 test/colortest.py | 22 ++++++++++----------
 2 files changed, 34 insertions(+), 39 deletions(-)

diff --git a/mmgen/color.py b/mmgen/color.py
index 34dcade8..d73ed66f 100755
--- a/mmgen/color.py
+++ b/mmgen/color.py
@@ -41,21 +41,6 @@ _colors = {
 	'yelbg':       ( (232,229), (30,103) ),
 }
 
-_reset = ''
-
-for _c in _colors:
-	_e = _colors[_c]
-	globals()['_256_'+_c] = (
-		'\033[38;5;{};1m'.format(_e[0]) if type(_e[0]) == int else
-		'\033[38;5;{};48;5;{};1m'.format(*_e[0])
-	)
-	globals()['_16_'+_c] = (
-		'\033[{}m'.format(_e[1][0]) if _e[1][1] == 0 else
-		'\033[{};{}m'.format(*_e[1])
-	)
-	globals()['_clr_'+_c] = ''
-	exec(f'{_c} = lambda s: _clr_{_c}+s+_reset')
-
 def nocolor(s):
 	return s
 
@@ -92,22 +77,30 @@ def get_terminfo_colors(term=None):
 def init_color(num_colors='auto'):
 	assert num_colors in ('auto',8,16,256,0)
 
-	globals()['_reset'] = '\033[0m' if num_colors else ''
-
-	if num_colors in (0,8,16):
-		pfx = '_16_'
-	else:
+	if num_colors == 'auto':
 		import os
 		t = os.getenv('TERM')
-		if num_colors == 256 or (t and t.endswith('256color')) or get_terminfo_colors() == 256:
-			pfx = '_256_'
-		else:
-			pfx = '_16_'
+		num_colors = 256 if (t and t.endswith('256color')) or get_terminfo_colors() == 256 else 16
 
-	for c in _colors:
-		if num_colors == 0:
-			globals()['_clr_'+c] = ''
-		else:
-			globals()['_clr_'+c] = globals()[pfx+c]
+	reset = '\033[0m'
+	if num_colors == 0:
+		ncc = (lambda s: s).__code__
+		for c in _colors:
+			globals()[c].__code__ = ncc
+	elif num_colors == 256:
+		for c,e in _colors.items():
+			start = (
+				'\033[38;5;{};1m'.format(e[0]) if type(e[0]) == int else
+				'\033[38;5;{};48;5;{};1m'.format(*e[0]) )
+			globals()[c].__code__ = eval(f'(lambda s: "{start}" + s + "{reset}").__code__')
+	elif num_colors in (8,16):
+		for c,e in _colors.items():
+			start = (
+				'\033[{}m'.format(e[1][0]) if e[1][1] == 0 else
+				'\033[{};{}m'.format(*e[1]) )
+			globals()[c].__code__ = eval(f'(lambda s: "{start}" + s + "{reset}").__code__')
 
 	set_vt100()
+
+for _c in _colors:
+	exec(f'{_c} = lambda s: s')
diff --git a/test/colortest.py b/test/colortest.py
index c3173afb..a71354b2 100755
--- a/test/colortest.py
+++ b/test/colortest.py
@@ -14,8 +14,17 @@ from mmgen.color import _colors
 
 def test_color():
 
+	ymsg("Terminal display:") # init_color() not called yet, so no yellow here
+
+	for desc,nc in (('pre-init',None),('auto','auto'),('8-color',8),('256-color',256),('disabled',0)):
+		if nc != None:
+			init_color(num_colors=nc)
+		msg('{:9}: {}'.format(
+			desc,
+			' '.join(globals()[c](c) for c in sorted(_colors)) ))
+
 	init_color()
-	gmsg("Parsed terminfo 'colors' values:")
+	gmsg("\nParsed terminfo 'colors' values:")
 
 	for t,c in (('rxvt',8),('xterm',8),('rxvt-unicode',88),('screen-256color',256),('xterm-256color',256)):
 		ret = get_terminfo_colors(t)
@@ -23,18 +32,11 @@ def test_color():
 			set_vt100()
 			ymsg(f'Warning: unable to get info for terminal {t!r}')
 			continue
-		msg(f'{t}: {ret}')
+		msg(f'{t}: {orange(str(ret))}')
 		assert c == ret, f"'colors' value for terminal {t} ({ret}) does not match expected value of {c}"
 
 	ret = get_terminfo_colors()
-	msg(f'This terminal ({os.getenv("TERM")}): {ret}')
+	msg(f'{os.getenv("TERM")} (this terminal): {orange(str(ret))}')
 	set_vt100()
-	gmsg("Terminal display:")
-
-	for desc,n in (('auto','auto'),('8-color',8),('256-color',256),('off',0)):
-		init_color(num_colors=n)
-		msg('{:9}: {}'.format(
-			desc,
-			' '.join(globals()[c](c) for c in sorted(_colors)) ))
 
 test_color()