From 5617d30e12c54b51744b58eb1c00f0bce6a93579 Mon Sep 17 00:00:00 2001
From: Renan Rodrigo <renanrodrigo@canonical.com>
Date: Thu, 12 Sep 2024 01:14:26 -0300
Subject: [PATCH] cli: add colorize method for the new CLI

This method should be the standard way of colorizing any user facing
string in the pro client, in the future.

Signed-off-by: Renan Rodrigo <renanrodrigo@canonical.com>
---
 uaclient/cli/formatter.py                | 23 +++++++++++++++++++++++
 uaclient/cli/tests/test_cli_formatter.py | 19 +++++++++++++++++++
 2 files changed, 42 insertions(+)

diff --git a/uaclient/cli/formatter.py b/uaclient/cli/formatter.py
index c32bb4e4e5..0b179cc7ed 100644
--- a/uaclient/cli/formatter.py
+++ b/uaclient/cli/formatter.py
@@ -1,10 +1,22 @@
 import os
 import re
 import sys
+from enum import Enum
 
 from uaclient.config import UAConfig
 
 
+class TxtColor(Enum):
+    BLUE = "\033[94m"
+    BOLD = "\033[1m"
+    GREEN = "\033[92m"
+    GREY = "\033[37m"
+    ORANGE = "\033[38;5;208m"
+    RED = "\033[91m"
+    YELLOW = "\033[93m"
+    ENDC = "\033[0m"
+
+
 # Class attributes and methods so we don't need singletons or globals for this
 class ProOutputFormatterConfig:
     use_utf8 = True
@@ -57,3 +69,14 @@ def decolorize(self):
     def split(self, sep=None, maxsplit=-1):
         splitted = super().split(sep=sep, maxsplit=maxsplit)
         return [self.__class__(s) for s in splitted]
+
+
+def colorize(input_string: str, color: TxtColor) -> str:
+    if not ProOutputFormatterConfig.use_color:
+        return input_string
+
+    return_string = ProColorString(color.value + input_string)
+    if not return_string.endswith(TxtColor.ENDC.value):
+        return_string += TxtColor.ENDC.value
+
+    return return_string
diff --git a/uaclient/cli/tests/test_cli_formatter.py b/uaclient/cli/tests/test_cli_formatter.py
index eacb8987d3..5cb7833da1 100644
--- a/uaclient/cli/tests/test_cli_formatter.py
+++ b/uaclient/cli/tests/test_cli_formatter.py
@@ -3,6 +3,7 @@
 
 from uaclient.cli.formatter import ProColorString
 from uaclient.cli.formatter import ProOutputFormatterConfig as POFC
+from uaclient.cli.formatter import TxtColor, colorize
 
 M_PATH = "uaclient.cli.formatter."
 
@@ -155,3 +156,21 @@ def test_split_produces_color_strings(self):
         ] == splitted
         for chunk in splitted:
             assert isinstance(chunk, ProColorString)
+
+
+class TestColorize:
+    def test_adds_color_to_strings(self, FakeConfig):
+        POFC.use_color = True
+        assert "\033[94mtest_string\033[0m" == colorize(
+            "test_string", TxtColor.BLUE
+        )
+
+    def test_noop_if_no_color_set(self, FakeConfig):
+        POFC.use_color = False
+        assert "test_string" == colorize("test_string", TxtColor.BLUE)
+
+    def does_not_add_multiple_end_characters(self, FakeConfig):
+        POFC.use_color = True
+        blue_string = colorize("test_string", TxtColor.BLUE)
+        bold_blue_string = colorize(blue_string, TxtColor.BOLD)
+        assert "\033[1m\033[94mtest_string\033[0m" == bold_blue_string