fix: switch to rich for color logging. (aws#5107)

* fix: switch to `rich` for color logging.

- maintain click backcompatability with interface.

* tests: add unit tests

- import the logger name instead of hardcoding it.

* fix: add Enums for colors

* lint: make lint happy

samcli/commands/sync/command.py

@@ -284,7 +284,7 @@ def do_cli(
             return
         global_config.set_accelerate_opt_in_stack(template_file, stack_name)
     else:
-        LOG.info(Colored().yellow(SYNC_INFO_TEXT))
+        LOG.info(Colored().color_log(msg=SYNC_INFO_TEXT, color="yellow"), extra=dict(markup=True))
 
     s3_bucket_name = s3_bucket or manage_stack(profile=profile, region=region)
 

samcli/lib/build/app_builder.py

@@ -35,7 +35,7 @@
 from samcli.lib.samlib.resource_metadata_normalizer import ResourceMetadataNormalizer
 from samcli.lib.docker.log_streamer import LogStreamer, LogStreamError
 from samcli.lib.providers.provider import ResourcesToBuildCollector, get_full_path, Stack
-from samcli.lib.utils.colors import Colored
+from samcli.lib.utils.colors import Colored, Colors
 from samcli.lib.utils import osutils
 from samcli.lib.utils.packagetype import IMAGE, ZIP
 from samcli.lib.utils.stream_writer import StreamWriter
@@ -641,7 +641,7 @@ def _build_function(  # pylint: disable=R1710
                     f"update to a newer supported runtime. For more information please check AWS Lambda Runtime "
                     f"Support Policy: https://docs.aws.amazon.com/lambda/latest/dg/runtime-support-policy.html"
                 )
-                LOG.warning(self._colored.yellow(message))
+                LOG.warning(self._colored.color_log(msg=message, color=Colors.WARNING), extra=dict(markup=True))
                 raise UnsupportedRuntimeException(f"Building functions with {runtime} is no longer supported")
 
             # Create the arguments to pass to the builder

samcli/lib/deploy/deployer.py

@@ -37,7 +37,7 @@
 from samcli.lib.deploy.utils import DeployColor, FailureMode
 from samcli.lib.package.local_files_utils import get_uploaded_s3_object_name, mktempfile
 from samcli.lib.package.s3_uploader import S3Uploader
-from samcli.lib.utils.colors import Colored
+from samcli.lib.utils.colors import Colored, Colors
 from samcli.lib.utils.time import utc_to_timestamp
 
 LOG = logging.getLogger(__name__)
@@ -510,7 +510,7 @@ def wait_for_execute(
             if disable_rollback and on_failure is not FailureMode.DELETE:
                 # This will only display the message if disable rollback is set or if DO_NOTHING is specified
                 msg = self._gen_deploy_failed_with_rollback_disabled_msg(stack_name)
-                LOG.info(self._colored.red(msg))
+                LOG.info(self._colored.color_log(msg=msg, color=Colors.FAILURE), extra=dict(markup=True))
 
             raise deploy_exceptions.DeployFailedError(stack_name=stack_name, msg=str(ex))
 
@@ -637,7 +637,7 @@ def sync(
                 self.wait_for_execute(stack_name, "CREATE", disable_rollback, on_failure=on_failure)
                 msg = "\nStack creation succeeded. Sync infra completed.\n"
 
-            LOG.info(self._colored.green(msg))
+            LOG.info(self._colored.color_log(msg=msg, color=Colors.SUCCESS), extra=dict(markup=True))
 
             return result
         except botocore.exceptions.ClientError as ex:
@@ -661,7 +661,7 @@ def _display_stack_outputs(stack_outputs: List[Dict], **kwargs) -> None:
                     format_string=OUTPUTS_FORMAT_STRING,
                     format_args=kwargs["format_args"],
                     columns_dict=OUTPUTS_DEFAULTS_ARGS.copy(),
-                    color="green",
+                    color=Colors.SUCCESS,
                     replace_whitespace=False,
                     break_long_words=False,
                     drop_whitespace=False,

samcli/lib/providers/sam_function_provider.py

@@ -9,7 +9,7 @@
 from samcli.commands._utils.template import TemplateFailedParsingException
 from samcli.commands.local.cli_common.user_exceptions import InvalidLayerVersionArn
 from samcli.lib.providers.exceptions import InvalidLayerReference
-from samcli.lib.utils.colors import Colored
+from samcli.lib.utils.colors import Colored, Colors
 from samcli.lib.utils.file_observer import FileObserver
 from samcli.lib.utils.packagetype import IMAGE, ZIP
 from samcli.lib.utils.resources import (
@@ -163,7 +163,7 @@ def _deprecate_notification(self, runtime: Optional[str]) -> None:
                 "runtime. For more information please check AWS Lambda Runtime Support Policy: "
                 "https://docs.aws.amazon.com/lambda/latest/dg/runtime-support-policy.html"
             )
-            LOG.warning(self._colored.yellow(message))
+            LOG.warning(self._colored.color_log(msg=message, color=Colors.WARNING), extra=dict(markup=True))
 
     def get_all(self) -> Iterator[Function]:
         """

samcli/lib/sync/sync_flow_executor.py

@@ -230,7 +230,10 @@ def _execute_step(
                 # Put it into deferred_tasks and add all of them at the end to avoid endless loop
                 if sync_flow_future:
                     self._running_futures.add(sync_flow_future)
-                    LOG.info(self._color.cyan(f"Syncing {sync_flow_future.sync_flow.log_name}..."))
+                    LOG.info(
+                        self._color.color_log(msg=f"Syncing {sync_flow_future.sync_flow.log_name}...", color="cyan"),
+                        extra=dict(markup=True),
+                    )
                 else:
                     deferred_tasks.append(sync_flow_task)
 
@@ -306,7 +309,10 @@ def _handle_result(
             sync_flow_result: SyncFlowResult = future.result()
             for dependent_sync_flow in sync_flow_result.dependent_sync_flows:
                 self.add_sync_flow(dependent_sync_flow)
-            LOG.info(self._color.green(f"Finished syncing {sync_flow_result.sync_flow.log_name}."))
+            LOG.info(
+                self._color.color_log(msg=f"Finished syncing {sync_flow_result.sync_flow.log_name}.", color="green"),
+                extra=dict(markup=True),
+            )
         return True
 
     @staticmethod

samcli/lib/sync/watch_manager.py

@@ -15,7 +15,7 @@
 from samcli.lib.sync.infra_sync_executor import InfraSyncExecutor, InfraSyncResult
 from samcli.lib.sync.sync_flow_factory import SyncFlowFactory
 from samcli.lib.utils.code_trigger_factory import CodeTriggerFactory
-from samcli.lib.utils.colors import Colored
+from samcli.lib.utils.colors import Colored, Colors
 from samcli.lib.utils.path_observer import HandlerObserver
 from samcli.lib.utils.resource_trigger import OnChangeCallback, TemplateTrigger
 
@@ -96,10 +96,12 @@ def queue_infra_sync(self) -> None:
         """
         if self._disable_infra_syncs:
             LOG.info(
-                self._color.yellow(
-                    "You have enabled the --code flag, which limits sam sync updates to code changes only. To do a "
-                    "complete infrastructure and code sync, remove the --code flag."
-                )
+                self._color.color_log(
+                    msg="You have enabled the --code flag, which limits sam sync updates to code changes only. To do a "
+                    "complete infrastructure and code sync, remove the --code flag.",
+                    color=Colors.WARNING,
+                ),
+                extra=dict(markup=True),
             )
             return
         self._waiting_infra_sync = True
@@ -131,8 +133,12 @@ def _add_code_triggers(self) -> None:
                 trigger = self._trigger_factory.create_trigger(resource_id, self._on_code_change_wrapper(resource_id))
             except (MissingCodeUri, MissingLocalDefinition):
                 LOG.warning(
-                    self._color.yellow("CodeTrigger not created as CodeUri or DefinitionUri is missing for %s."),
+                    self._color.color_log(
+                        msg="CodeTrigger not created as CodeUri or DefinitionUri is missing for %s.",
+                        color=Colors.WARNING,
+                    ),
                     str(resource_id),
+                    extra=dict(markup=True),
                 )
                 continue
 
@@ -149,7 +155,12 @@ def _add_template_triggers(self) -> None:
             try:
                 template_trigger.validate_template()
             except InvalidTemplateFile:
-                LOG.warning(self._color.yellow("Template validation failed for %s in %s"), template, stack.name)
+                LOG.warning(
+                    self._color.color_log(msg="Template validation failed for %s in %s", color=Colors.WARNING),
+                    template,
+                    stack.name,
+                    extra=dict(markup=True),
+                )
 
             self._observer.schedule_handlers(template_trigger.get_path_handlers())
 
@@ -192,13 +203,17 @@ def start(self) -> None:
             self.queue_infra_sync()
             if self._disable_infra_syncs:
                 self._start_sync()
-                LOG.info(self._color.green("Sync watch started."))
+                LOG.info(
+                    self._color.color_log(msg="Sync watch started.", color=Colors.SUCCESS), extra=dict(markup=True)
+                )
             self._start()
         except KeyboardInterrupt:
-            LOG.info(self._color.cyan("Shutting down sync watch..."))
+            LOG.info(
+                self._color.color_log(msg="Shutting down sync watch...", color=Colors.PROGRESS), extra=dict(markup=True)
+            )
             self._observer.stop()
             self._stop_code_sync()
-            LOG.info(self._color.green("Sync watch stopped."))
+            LOG.info(self._color.color_log(msg="Sync watch stopped.", color=Colors.SUCCESS), extra=dict(markup=True))
 
     def _start(self) -> None:
         """Start WatchManager and watch for changes to the template and its code resources."""
@@ -220,16 +235,24 @@ def _start_sync(self) -> None:
 
     def _execute_infra_sync(self, first_sync: bool = False) -> None:
         """Logic to execute infra sync."""
-        LOG.info(self._color.cyan("Queued infra sync. Waiting for in progress code syncs to complete..."))
+        LOG.info(
+            self._color.color_log(
+                msg="Queued infra sync. Waiting for in progress code syncs to complete...", color=Colors.PROGRESS
+            ),
+            extra=dict(markup=True),
+        )
         self._waiting_infra_sync = False
         self._stop_code_sync()
         try:
-            LOG.info(self._color.cyan("Starting infra sync."))
+            LOG.info(self._color.color_log(msg="Starting infra sync.", color=Colors.PROGRESS), extra=dict(markup=True))
             infra_sync_result = self._execute_infra_context(first_sync)
         except Exception as e:
             LOG.error(
-                self._color.red("Failed to sync infra. Code sync is paused until template/stack is fixed."),
+                self._color.color_log(
+                    msg="Failed to sync infra. Code sync is paused until template/stack is fixed.", color=Colors.FAILURE
+                ),
                 exc_info=e,
+                extra=dict(markup=True),
             )
             # Unschedule all triggers and only add back the template one as infra sync is incorrect.
             self._observer.unschedule_all()
@@ -245,12 +268,18 @@ def _execute_infra_sync(self, first_sync: bool = False) -> None:
                 # To improve: only initiate code syncs for ones with template changes
                 self._queue_up_code_syncs(infra_sync_result.code_sync_resources)
                 LOG.info(
-                    self._color.green("Skipped infra sync as the local template is in sync with the cloud template.")
+                    self._color.color_log(
+                        msg="Skipped infra sync as the local template is in sync with the cloud template.",
+                        color=Colors.SUCCESS,
+                    ),
+                    extra=dict(markup=True),
                 )
                 if len(infra_sync_result.code_sync_resources) != 0:
                     LOG.info("Required code syncs are queued up.")
             else:
-                LOG.info(self._color.green("Infra sync completed."))
+                LOG.info(
+                    self._color.color_log(msg="Infra sync completed.", color=Colors.SUCCESS), extra=dict(markup=True)
+                )
 
     def _queue_up_code_syncs(self, resource_ids_with_code_sync: Set[ResourceIdentifier]) -> None:
         """
@@ -300,15 +329,25 @@ def _watch_sync_flow_exception_handler(self, sync_flow_exception: SyncFlowExcept
         """
         exception = sync_flow_exception.exception
         if isinstance(exception, MissingPhysicalResourceError):
-            LOG.warning(self._color.yellow("Missing physical resource. Infra sync will be started."))
+            LOG.warning(
+                self._color.color_log(
+                    msg="Missing physical resource. Infra sync will be started.", color=Colors.WARNING
+                ),
+                extra=dict(markup=True),
+            )
             self.queue_infra_sync()
         elif isinstance(exception, InfraSyncRequiredError):
             LOG.warning(
                 self._color.yellow(
                     f"Infra sync is required for {exception.resource_identifier} due to: "
                     + f"{exception.reason}. Infra sync will be started."
-                )
+                ),
+                extra=dict(markup=True),
             )
             self.queue_infra_sync()
         else:
-            LOG.error(self._color.red("Code sync encountered an error."), exc_info=exception)
+            LOG.error(
+                self._color.color_log(msg="Code sync encountered an error.", color=Colors.FAILURE),
+                exc_info=exception,
+                extra=dict(markup=True),
+            )

samcli/lib/utils/colors.py

@@ -2,10 +2,17 @@
 Wrapper to generated colored messages for printing in Terminal
 """
 
+import logging
 import os
 import platform
+from enum import Enum
 
 import click
+from rich.logging import RichHandler
+from rich.style import Style
+from rich.text import Text
+
+from samcli.lib.utils.sam_logging import SAM_CLI_LOGGER_NAME
 
 # Enables ANSI escape codes on Windows
 if platform.system().lower() == "windows":
@@ -15,6 +22,14 @@
         pass
 
 
+# Python 3.11 has StrEnum
+class Colors(str, Enum):
+    SUCCESS = "green"
+    FAILURE = "red"
+    WARNING = "yellow"
+    PROGRESS = "cyan"
+
+
 class Colored:
     """
     Helper class to add ANSI colors and decorations to text. Given a string, ANSI colors are added with special prefix
@@ -42,6 +57,9 @@ def __init__(self, colorize=True):
         colorize : bool
             Optional. Set this to True to turn on coloring. False will turn off coloring
         """
+        self.rich_logging = any(
+            isinstance(handler, RichHandler) for handler in logging.getLogger(SAM_CLI_LOGGER_NAME).handlers
+        )
         self.colorize = colorize
 
     def red(self, msg):
@@ -68,6 +86,14 @@ def underline(self, msg):
         """Underline the input"""
         return click.style(msg, underline=True) if self.colorize else msg
 
+    def underline_log(self, msg):
+        """Underline the input such that underlying Rich Logger understands it (if configured)."""
+        if self.rich_logging:
+            _color_msg = Text(msg, style=Style(underline=True))
+            return _color_msg.markup if self.colorize else msg
+        else:
+            return click.style(msg, underline=True) if self.colorize else msg
+
     def bold(self, msg):
         """Bold the input"""
         return click.style(msg, bold=True) if self.colorize else msg
@@ -76,3 +102,12 @@ def _color(self, msg, color):
         """Internal helper method to add colors to input"""
         kwargs = {"fg": color}
         return click.style(msg, **kwargs) if self.colorize else msg
+
+    def _color_log(self, msg, color):
+        """Marked up text with color used for logging with a logger"""
+        _color_msg = Text(msg, style=Style(color=color))
+        return _color_msg.markup if self.colorize else msg
+
+    def color_log(self, msg, color):
+        """Internal helper method to add colors such that underlying Rich Logger understands it (if configured)."""
+        return self._color_log(msg=msg, color=color) if self.rich_logging else self._color(msg=msg, color=color)

tests/unit/lib/utils/test_colors.py

@@ -1,12 +1,17 @@
+from logging import StreamHandler
 from unittest import TestCase
+from unittest.mock import patch
+
 from parameterized import parameterized, param
+from rich.logging import RichHandler
 
 from samcli.lib.utils.colors import Colored
 
 
 class TestColored(TestCase):
     def setUp(self):
         self.msg = "message"
+        self.color = "green"
 
     @parameterized.expand(
         [
@@ -26,3 +31,15 @@ def test_various_decorations(self, decoration_name, ansi_prefix):
 
         self.assertEqual(expected, getattr(with_color, decoration_name)(self.msg))
         self.assertEqual(self.msg, getattr(without_color, decoration_name)(self.msg))
+
+    @parameterized.expand(
+        [
+            param([RichHandler()], "[green]message[/green]"),
+            param([StreamHandler()], "\x1b[32mmessage\x1b[0m"),
+        ]
+    )
+    @patch("samcli.lib.utils.colors.logging")
+    def test_color_log_log_handlers(self, log_handler, response, mock_logger):
+        mock_logger.getLogger().handlers = log_handler
+        with_rich_color = Colored(colorize=True)
+        self.assertEqual(with_rich_color.color_log(msg=self.msg, color=self.color), response)