terminalwriter: factor out pygments lexer & formatter selection to own functions

We intend to add some more logic here.

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,38 +200,68 @@ class TerminalWriter:
         for indent, new_line in zip(indents, new_lines):
             self.line(indent + new_line)
 
-    def _highlight(
+    def _get_pygments_lexer(
-        self, source: str, lexer: Literal["diff", "python"] = "python"
+        self, lexer: Literal["python", "diff"]
-    ) -> str:
+    ) -> Optional["Lexer"]:
-        """Highlight the given source if we have markup support."""
+        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:
+        theme = os.getenv("PYTEST_THEME")
-            return source
+        theme_mode = os.getenv("PYTEST_THEME_MODE", "dark")
 
         try:
             from pygments.formatters.terminal import TerminalFormatter
 
-            if lexer == "python":
+            return TerminalFormatter(bg=theme_mode, style=theme)
-                from pygments.lexers.python import PythonLexer as Lexer
+
-            elif lexer == "diff":
+        except pygments.util.ClassNotFound as e:
-                from pygments.lexers.diff import DiffLexer as Lexer
+            raise UsageError(
-            else:
+                f"PYTEST_THEME environment variable had an invalid value: '{theme}'. "
-                assert_never(lexer)
+                "Only valid pygment styles are allowed."
-            from pygments import highlight
+            ) from e
-            import pygments.util
+        except pygments.util.OptionError as e:
-        except ImportError:
+            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
-        else:
+
-            try:
+        pygments_lexer = self._get_pygments_lexer(lexer)
-                highlighted: str = highlight(
+        if pygments_lexer is None:
-                    source,
+            return source
-                    Lexer(),
+
-                    TerminalFormatter(
+        pygments_formatter = self._get_pygments_formatter()
-                        bg=os.getenv("PYTEST_THEME_MODE", "dark"),
+        if pygments_formatter is None:
-                        style=os.getenv("PYTEST_THEME"),
+            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":
@@ -234,18 +270,6 @@ class TerminalWriter:
         # 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
+        highlighted = "\x1b[0m" + highlighted
-            except pygments.util.ClassNotFound as e:
+
-                raise UsageError(
+        return highlighted
-                    "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