[3.15] gh-154335: Allow disabling terminal colors in Tachyon's pstats_collector module (GH-154344) (#158122)

gh-154335: Allow disabling terminal colors in Tachyon's `pstats_collector` module (GH-154344)

* Use _colorize.get_colors() to allow disabling ANSI escape codes through environment variables

* add news entry

* use get_colors() instead of ANSIColors enum also in Lib/profiling/sampling/sample.py

* fix data_lines collection of pstats lines for no-color mode

* preserve lazy import of _colorize

* gh-154335: Cover profiler color policy and update NEWS

---------
(cherry picked from commit 0546b5f)

Co-authored-by: Peter Fackeldey <fackeldey.peter@gmail.com>
Co-authored-by: Pablo Galindo Salgado <pablogsal@gmail.com>

Author: 3 people

Lib/profiling/sampling/pstats_collector.py

@@ -1,7 +1,7 @@
 import collections
 import marshal
 import pstats
-lazy from _colorize import ANSIColors
+lazy from _colorize import get_colors
 
 from .collector import Collector, extract_lineno
 from .constants import MICROSECONDS_PER_SECOND, PROFILING_MODE_CPU
@@ -178,6 +178,8 @@ def print_stats(self, sort=-1, limit=None, show_summary=True, mode=None):
         }
 
         # Print header with colors and proper alignment
+        ANSIColors = get_colors()
+
         print(f"{ANSIColors.BOLD_BLUE}Profile Stats:{ANSIColors.RESET}")
 
         header_nsamples = f"{ANSIColors.BOLD_BLUE}{'nsamples':>{col_widths['nsamples']}}{ANSIColors.RESET}"
@@ -269,6 +271,8 @@ def _determine_best_unit(max_value):
 
     def _print_summary(self, stats_list, total_samples):
         """Print summary of interesting functions."""
+        ANSIColors = get_colors()
+
         print(
             f"\n{ANSIColors.BOLD_BLUE}Summary of Interesting Functions:{ANSIColors.RESET}"
         )

Lib/profiling/sampling/sample.py

@@ -6,7 +6,7 @@
 import sysconfig
 import time
 from collections import deque
-lazy from _colorize import ANSIColors
+lazy from _colorize import get_colors
 
 from .pstats_collector import PstatsCollector
 from .stack_collector import CollapsedStackCollector, FlamegraphCollector
@@ -276,6 +276,8 @@ def _print_realtime_stats(self):
         )  # Max time = Min Hz
 
         # Build cache stats string if stats collection is enabled
+        ANSIColors = get_colors()
+
         cache_stats_str = ""
         if self.collect_stats:
             try:
@@ -309,6 +311,8 @@ def _print_unwinder_stats(self):
         except RuntimeError:
             return  # Stats not enabled
 
+        ANSIColors = get_colors()
+
         print(f"\n{ANSIColors.BOLD_BLUE}{'='*50}{ANSIColors.RESET}")
         print(f"{ANSIColors.BOLD_BLUE}Unwinder Statistics:{ANSIColors.RESET}")
 
@@ -403,6 +407,8 @@ def _print_binary_stats(self, collector):
         except (ValueError, RuntimeError):
             return  # Collector closed or stats unavailable
 
+        ANSIColors = get_colors()
+
         print(f"  {ANSIColors.CYAN}Binary Encoding:{ANSIColors.RESET}")
 
         repeat_records = stats.get('repeat_records', 0)

Lib/test/test_profiling/test_sampling_profiler/test_profiler.py

@@ -17,7 +17,7 @@
         "Test only runs when _remote_debugging is available"
     )
 
-from test.support import force_not_colorized_test_class
+from test.support import force_colorized, force_not_colorized_test_class
 
 
 def print_sampled_stats(stats, sort=-1, limit=None, show_summary=True, sample_interval_usec=100):
@@ -484,6 +484,15 @@ def test_print_sampled_stats_basic(self):
         self.assertIn("func1", result)
         self.assertIn("func2", result)
         self.assertIn("func3", result)
+        self.assertNotIn("\x1b[", result)
+
+    @force_colorized
+    def test_print_sampled_stats_colorized(self):
+        with io.StringIO() as output, mock.patch("sys.stdout", output):
+            print_sampled_stats(self.mock_stats)
+            result = output.getvalue()
+
+        self.assertIn("\x1b[1;34mProfile Stats:", result)
 
     def test_print_sampled_stats_sorting(self):
         """Test different sorting options."""
@@ -751,6 +760,7 @@ def test_print_sampled_stats_sort_by_name(self):
                 and not "calls" in line  # Skip summary lines
                 and not "total time" in line  # Skip summary lines
                 and not "cumulative time" in line
+                and not "filename:lineno(function)" in line  # Skip header line
             ):  # Skip summary lines
                 data_lines.append(line)
 

Misc/NEWS.d/next/Library/2026-07-21-12-23-48.gh-issue-154335.SRf8Gr.rst

@@ -0,0 +1,3 @@
+The :mod:`profiling.sampling` profiler now respects ``NO_COLOR``,
+``FORCE_COLOR``, and ``PYTHON_COLORS`` and avoids colored output when
+not writing to a terminal.