From b82dfd174591f58868e29f8838d18f371632dba5 Mon Sep 17 00:00:00 2001
From: Pradyun Gedam <pradyunsg@users.noreply.github.com>
Date: Sat, 4 Dec 2021 18:47:54 +0000
Subject: [PATCH] Restore colors in regular logging messages

Also, reflects that `--no-color` does not mean no-ansi-escapes but only
affects colour rules.
---
 src/pip/_internal/utils/logging.py | 14 +++++++++++---
 tests/functional/test_no_color.py  | 22 +++++++++++-----------
 2 files changed, 22 insertions(+), 14 deletions(-)

diff --git a/src/pip/_internal/utils/logging.py b/src/pip/_internal/utils/logging.py
index 8a71535dc..0d505dc87 100644
--- a/src/pip/_internal/utils/logging.py
+++ b/src/pip/_internal/utils/logging.py
@@ -18,6 +18,7 @@ from pip._vendor.rich.console import (
 from pip._vendor.rich.highlighter import NullHighlighter
 from pip._vendor.rich.logging import RichHandler
 from pip._vendor.rich.segment import Segment
+from pip._vendor.rich.style import Style
 
 from pip._internal.exceptions import DiagnosticPipError
 from pip._internal.utils._log import VERBOSE, getLogger
@@ -148,6 +149,8 @@ class RichPipStreamHandler(RichHandler):
 
     # Our custom override on rich's logger, to make things work as we need them to.
     def emit(self, record: logging.LogRecord) -> None:
+        style: Optional[Style] = None
+
         # If we are given a diagnostic error to present, present it with indentation.
         if record.msg == "[present-diagnostic]" and len(record.args) == 1:
             diagnostic_error: DiagnosticPipError = record.args[0]  # type: ignore[index]
@@ -159,9 +162,14 @@ class RichPipStreamHandler(RichHandler):
         else:
             message = self.format(record)
             renderable = self.render_message(record, message)
+            if record.levelno is not None:
+                if record.levelno >= logging.ERROR:
+                    style = Style(color="red")
+                elif record.levelno >= logging.WARNING:
+                    style = Style(color="yellow")
 
         try:
-            self.console.print(renderable, overflow="ignore", crop=False)
+            self.console.print(renderable, overflow="ignore", crop=False, style=style)
         except Exception:
             self.handleError(record)
 
@@ -252,7 +260,6 @@ def setup_logging(verbosity: int, no_color: bool, user_log_file: Optional[str])
         "stderr": "ext://sys.stderr",
     }
     handler_classes = {
-        "subprocess": "logging.StreamHandler",
         "stream": "pip._internal.utils.logging.RichPipStreamHandler",
         "file": "pip._internal.utils.logging.BetterRotatingFileHandler",
     }
@@ -310,8 +317,9 @@ def setup_logging(verbosity: int, no_color: bool, user_log_file: Optional[str])
                 # from the "subprocessor" logger.
                 "console_subprocess": {
                     "level": level,
-                    "class": handler_classes["subprocess"],
+                    "class": handler_classes["stream"],
                     "stream": log_streams["stderr"],
+                    "no_color": no_color,
                     "filters": ["restrict_to_subprocess"],
                     "formatter": "indent",
                 },
diff --git a/tests/functional/test_no_color.py b/tests/functional/test_no_color.py
index 9ead9996a..4094bdd16 100644
--- a/tests/functional/test_no_color.py
+++ b/tests/functional/test_no_color.py
@@ -2,13 +2,16 @@
 Test specific for the --no-color option
 """
 import os
+import shutil
 import subprocess
+import sys
 
 import pytest
 
 from tests.lib import PipTestEnvironment
 
 
+@pytest.mark.skipif(shutil.which("script") is None, reason="no 'script' executable")
 def test_no_color(script: PipTestEnvironment) -> None:
     """Ensure colour output disabled when --no-color is passed."""
     # Using 'script' in this test allows for transparently testing pip's output
@@ -19,12 +22,13 @@ def test_no_color(script: PipTestEnvironment) -> None:
     # 'script' and well as the mere use of the same.
     #
     # This test will stay until someone has the time to rewrite it.
-    command = (
-        "script --flush --quiet --return /tmp/pip-test-no-color.txt "
-        '--command "pip uninstall {} noSuchPackage"'
-    )
+    pip_command = "pip uninstall {} noSuchPackage"
+    if sys.platform == "darwin":
+        command = f"script -q /tmp/pip-test-no-color.txt {pip_command}"
+    else:
+        command = f'script -q /tmp/pip-test-no-color.txt --command "{pip_command}"'
 
-    def get_run_output(option: str) -> str:
+    def get_run_output(option: str = "") -> str:
         cmd = command.format(option)
         proc = subprocess.Popen(
             cmd,
@@ -33,8 +37,6 @@ def test_no_color(script: PipTestEnvironment) -> None:
             stderr=subprocess.PIPE,
         )
         proc.communicate()
-        if proc.returncode:
-            pytest.skip("Unable to capture output using script: " + cmd)
 
         try:
             with open("/tmp/pip-test-no-color.txt") as output_file:
@@ -43,7 +45,5 @@ def test_no_color(script: PipTestEnvironment) -> None:
         finally:
             os.unlink("/tmp/pip-test-no-color.txt")
 
-    assert "\x1b" in get_run_output(option=""), "Expected color in output"
-    assert "\x1b" not in get_run_output(
-        option="--no-color"
-    ), "Expected no color in output"
+    assert "\x1b[3" in get_run_output(""), "Expected color in output"
+    assert "\x1b[3" not in get_run_output("--no-color"), "Expected no color in output"