From 5003ba6b55133393159d6b46bc3778e75addf698 Mon Sep 17 00:00:00 2001 From: Montvydas Klumbys Date: Tue, 26 Mar 2019 13:59:15 +0000 Subject: [PATCH] Remove ESC char printing when not needed For "default" or unknown (e.g. None) colour do not print ESC characters. This is Useful when storing output into a text file or when running on windows cmd. --- bashplotlib/utils/helpers.py | 16 +++++++++------- 1 file changed, 9 insertions(+), 7 deletions(-) diff --git a/bashplotlib/utils/helpers.py b/bashplotlib/utils/helpers.py index cf209ee..8b37cf5 100644 --- a/bashplotlib/utils/helpers.py +++ b/bashplotlib/utils/helpers.py @@ -26,22 +26,24 @@ colour_help = ', '.join([colour for colour in bcolours if colour != "ENDC"]) -def get_colour(colour): +def get_colour(colour, default="default"): """ Get the escape code sequence for a colour """ - return bcolours.get(colour, bcolours['ENDC']) + return bcolours.get(colour, bcolours[default]) -def printcolour(text, sameline=False, colour=get_colour("ENDC")): +def printcolour(text, sameline=False, colour="default"): """ Print color text using escape codes """ - if sameline: - sep = '' + sep = '' if sameline else '\n' + + # If no colour set, do not print color ESC characters + if get_colour(colour) == get_colour("ENDC"): + sys.stdout.write(text + sep) else: - sep = '\n' - sys.stdout.write(get_colour(colour) + text + bcolours["ENDC"] + sep) + sys.stdout.write(get_colour(colour) + text + get_colour("ENDC") + sep) def drange(start, stop, step=1.0, include_stop=False):