From 24d2ec0bbaaed46f5e681a25e0964446fe093528 Mon Sep 17 00:00:00 2001
From: The MMGen Project <mmgen@tuta.io>
Date: Fri, 5 May 2023 10:58:44 +0000
Subject: [PATCH] devtools.pmsg(): add color support, add test

---
 mmgen/devtools.py                | 48 ++++++++++++++++++++------------
 test/unit_tests_d/ut_devtools.py | 20 ++++++++++++-
 2 files changed, 49 insertions(+), 19 deletions(-)

diff --git a/mmgen/devtools.py b/mmgen/devtools.py
index 185c09e5..a3001b3c 100755
--- a/mmgen/devtools.py
+++ b/mmgen/devtools.py
@@ -14,29 +14,40 @@ devtools: Developer tools for the MMGen suite
 
 import sys
 
-def pfmt(*args):
+color_funcs = {
+	name: lambda s, n=n: f'\033[{n};1m{s}\033[0m'
+		for name,n in (
+			('red',   31),
+			('green', 32),
+			('yellow',33),
+			('blue',  34),
+			('purple',35))
+}
+
+def pfmt(*args,color=None):
 	import pprint
-	return pprint.PrettyPrinter(indent=4).pformat(
+	ret = pprint.PrettyPrinter(indent=4).pformat(
 		args if len(args) > 1 else '' if not args else args[0] )
+	return color_funcs[color](ret) if color else ret
 
-def pmsg(*args):
-	sys.stderr.write(pfmt(*args) + '\n')
+def pmsg(*args,color=None):
+	sys.stderr.write(pfmt(*args,color=color) + '\n')
 
-def pmsg_r(*args):
-	sys.stderr.write(pfmt(*args))
+def pmsg_r(*args,color=None):
+	sys.stderr.write(pfmt(*args,color=color))
 
 def pdie(*args,exit_val=1):
-	pmsg(*args)
+	pmsg(*args,color='red')
 	sys.exit(exit_val)
 
 def pexit(*args):
 	pdie(*args,exit_val=0)
 
-def Pmsg(*args):
-	sys.stdout.write(pfmt(*args) + '\n')
+def Pmsg(*args,color=None):
+	sys.stdout.write(pfmt(*args,color=color) + '\n')
 
 def Pdie(*args,exit_val=1):
-	Pmsg(*args)
+	Pmsg(*args,color='red')
 	sys.exit(exit_val)
 
 def Pexit(*args):
@@ -98,17 +109,17 @@ def get_ndiff(a,b):
 class MMGenObjectMethods: # mixin class for MMGenObject
 
 	# Pretty-print an MMGenObject instance, recursing into sub-objects - WIP
-	def pmsg(self):
-		sys.stdout.write('\n'+self.pfmt())
+	def pmsg(self,color=None):
+		sys.stdout.write('\n'+self.pfmt(color=color))
 
-	def pdie(self,exit_val=1):
-		self.pmsg()
+	def pdie(self,exit_val=1,color=None):
+		self.pmsg(color=color)
 		sys.exit(exit_val)
 
-	def pexit(self):
-		self.pdie(exit_val=0)
+	def pexit(self,color=None):
+		self.pdie(exit_val=0,color=color)
 
-	def pfmt(self,lvl=0,id_list=[]):
+	def pfmt(self,lvl=0,id_list=[],color=None):
 		from decimal import Decimal
 		scalars = (str,int,float,Decimal)
 		def do_list(out,e,lvl=0,is_dict=False):
@@ -178,4 +189,5 @@ class MMGenObjectMethods: # mixin class for MMGenObject
 					l=(lvl*8)+4 ))
 
 		import re
-		return re.sub('\n+','\n',''.join(out))
+		ret = re.sub('\n+','\n',''.join(out))
+		return color_funcs[color](ret) if color else ret
diff --git a/test/unit_tests_d/ut_devtools.py b/test/unit_tests_d/ut_devtools.py
index d2ff528e..72e8c8a0 100755
--- a/test/unit_tests_d/ut_devtools.py
+++ b/test/unit_tests_d/ut_devtools.py
@@ -6,6 +6,7 @@ test.unit_tests_d.ut_devtools: devtools unit tests for the MMGen suite
 
 import os,json
 from mmgen.devtools import *
+from mmgen.util import msg
 from . import unit_tests_base
 
 textA = """
@@ -86,5 +87,22 @@ class unit_tests(unit_tests_base):
 			seed        = Seed(cfg,seed_bin=bytes.fromhex('bead'*16)),
 			addr_idxs   = '1',
 			mmtype      = 'B',
-			skip_chksum = True ).pmsg()
+			skip_chksum = True ).pmsg(color='green')
+		return True
+
+	def pmsg(self,name,ut):
+		colors = (None,'red','green','yellow','blue','purple')
+
+		msg('\npmsg_r():')
+		for color in colors:
+			pmsg_r({'color':color},color=color)
+
+		msg('\n\npmsg():')
+		for color in colors:
+			pmsg({'color':color},color=color)
+
+		msg('\nPmsg():')
+		for color in colors:
+			Pmsg({'color':color},color=color)
+
 		return True