From 67a570aea9e16ecf2694cea86bec326ce9056f9f Mon Sep 17 00:00:00 2001
From: Ran Benita <ran@unusedvar.com>
Date: Mon, 10 Jun 2024 11:38:33 +0300
Subject: [PATCH] terminalwriter: factor out pygments lexer & formatter
 selection to own functions

We intend to add some more logic here.
---
 src/_pytest/_io/terminalwriter.py | 120 ++++++++++++++++++------------
 1 file changed, 72 insertions(+), 48 deletions(-)

diff --git a/src/_pytest/_io/terminalwriter.py b/src/_pytest/_io/terminalwriter.py
index 5bcd05927..26344fd31 100644
--- a/src/_pytest/_io/terminalwriter.py
+++ b/src/_pytest/_io/terminalwriter.py
@@ -8,11 +8,17 @@ from typing import Literal
 from typing import Optional
 from typing import Sequence
 from typing import TextIO
+from typing import TYPE_CHECKING
 
 from ..compat import assert_never
 from .wcwidth import wcswidth
 
 
+if TYPE_CHECKING:
+    from pygments.formatter import Formatter
+    from pygments.lexer import Lexer
+
+
 # This code was initially copied from py 1.8.1, file _io/terminalwriter.py.
 
 
@@ -194,58 +200,76 @@ class TerminalWriter:
         for indent, new_line in zip(indents, new_lines):
             self.line(indent + new_line)
 
-    def _highlight(
-        self, source: str, lexer: Literal["diff", "python"] = "python"
-    ) -> str:
-        """Highlight the given source if we have markup support."""
+    def _get_pygments_lexer(
+        self, lexer: Literal["python", "diff"]
+    ) -> Optional["Lexer"]:
+        try:
+            if lexer == "python":
+                from pygments.lexers.python import PythonLexer
+
+                return PythonLexer()
+            elif lexer == "diff":
+                from pygments.lexers.diff import DiffLexer
+
+                return DiffLexer()
+            else:
+                assert_never(lexer)
+        except ModuleNotFoundError:
+            return None
+
+    def _get_pygments_formatter(self) -> Optional["Formatter"]:
+        try:
+            import pygments.util
+        except ModuleNotFoundError:
+            return None
+
         from _pytest.config.exceptions import UsageError
 
-        if not source or not self.hasmarkup or not self.code_highlight:
-            return source
+        theme = os.getenv("PYTEST_THEME")
+        theme_mode = os.getenv("PYTEST_THEME_MODE", "dark")
 
         try:
             from pygments.formatters.terminal import TerminalFormatter
 
-            if lexer == "python":
-                from pygments.lexers.python import PythonLexer as Lexer
-            elif lexer == "diff":
-                from pygments.lexers.diff import DiffLexer as Lexer
-            else:
-                assert_never(lexer)
-            from pygments import highlight
-            import pygments.util
-        except ImportError:
-            return source
-        else:
-            try:
-                highlighted: str = highlight(
-                    source,
-                    Lexer(),
-                    TerminalFormatter(
-                        bg=os.getenv("PYTEST_THEME_MODE", "dark"),
-                        style=os.getenv("PYTEST_THEME"),
-                    ),
-                )
-                # pygments terminal formatter may add a newline when there wasn't one.
-                # We don't want this, remove.
-                if highlighted[-1] == "\n" and source[-1] != "\n":
-                    highlighted = highlighted[:-1]
+            return TerminalFormatter(bg=theme_mode, style=theme)
 
-                # Some lexers will not set the initial color explicitly
-                # which may lead to the previous color being propagated to the
-                # start of the expression, so reset first.
-                return "\x1b[0m" + highlighted
-            except pygments.util.ClassNotFound as e:
-                raise UsageError(
-                    "PYTEST_THEME environment variable had an invalid value: '{}'. "
-                    "Only valid pygment styles are allowed.".format(
-                        os.getenv("PYTEST_THEME")
-                    )
-                ) from e
-            except pygments.util.OptionError as e:
-                raise UsageError(
-                    "PYTEST_THEME_MODE environment variable had an invalid value: '{}'. "
-                    "The only allowed values are 'dark' and 'light'.".format(
-                        os.getenv("PYTEST_THEME_MODE")
-                    )
-                ) from e
+        except pygments.util.ClassNotFound as e:
+            raise UsageError(
+                f"PYTEST_THEME environment variable had an invalid value: '{theme}'. "
+                "Only valid pygment styles are allowed."
+            ) from e
+        except pygments.util.OptionError as e:
+            raise UsageError(
+                f"PYTEST_THEME_MODE environment variable had an invalid value: '{theme_mode}'. "
+                "The only allowed values are 'dark' and 'light'."
+            ) from e
+
+    def _highlight(
+        self, source: str, lexer: Literal["diff", "python"] = "python"
+    ) -> str:
+        """Highlight the given source if we have markup support."""
+        if not source or not self.hasmarkup or not self.code_highlight:
+            return source
+
+        pygments_lexer = self._get_pygments_lexer(lexer)
+        if pygments_lexer is None:
+            return source
+
+        pygments_formatter = self._get_pygments_formatter()
+        if pygments_formatter is None:
+            return source
+
+        from pygments import highlight
+
+        highlighted: str = highlight(source, pygments_lexer, pygments_formatter)
+        # pygments terminal formatter may add a newline when there wasn't one.
+        # We don't want this, remove.
+        if highlighted[-1] == "\n" and source[-1] != "\n":
+            highlighted = highlighted[:-1]
+
+        # Some lexers will not set the initial color explicitly
+        # which may lead to the previous color being propagated to the
+        # start of the expression, so reset first.
+        highlighted = "\x1b[0m" + highlighted
+
+        return highlighted