Skip to content

Commit 0546b5f

Browse files
gh-154335: Allow disabling terminal colors in Tachyon's pstats_collector module (#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 --------- Co-authored-by: Pablo Galindo Salgado <pablogsal@gmail.com>
1 parent dbac447 commit 0546b5f

4 files changed

Lines changed: 26 additions & 3 deletions

File tree

‎Lib/profiling/sampling/pstats_collector.py‎

Lines changed: 5 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1,7 +1,7 @@
11
import collections
22
import marshal
33
import pstats
4-
lazy from _colorize import ANSIColors
4+
lazy from _colorize import get_colors
55

66
from .collector import Collector, extract_lineno
77
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):
178178
}
179179

180180
# Print header with colors and proper alignment
181+
ANSIColors = get_colors()
182+
181183
print(f"{ANSIColors.BOLD_BLUE}Profile Stats:{ANSIColors.RESET}")
182184

183185
header_nsamples = f"{ANSIColors.BOLD_BLUE}{'nsamples':>{col_widths['nsamples']}}{ANSIColors.RESET}"
@@ -269,6 +271,8 @@ def _determine_best_unit(max_value):
269271

270272
def _print_summary(self, stats_list, total_samples):
271273
"""Print summary of interesting functions."""
274+
ANSIColors = get_colors()
275+
272276
print(
273277
f"\n{ANSIColors.BOLD_BLUE}Summary of Interesting Functions:{ANSIColors.RESET}"
274278
)

‎Lib/profiling/sampling/sample.py‎

Lines changed: 7 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -6,7 +6,7 @@
66
import sysconfig
77
import time
88
from collections import deque
9-
lazy from _colorize import ANSIColors
9+
lazy from _colorize import get_colors
1010

1111
from .binary_collector import BinaryCollector
1212

@@ -272,6 +272,8 @@ def _print_realtime_stats(self):
272272
) # Max time = Min Hz
273273

274274
# Build cache stats string if stats collection is enabled
275+
ANSIColors = get_colors()
276+
275277
cache_stats_str = ""
276278
if self.collect_stats:
277279
try:
@@ -305,6 +307,8 @@ def _print_unwinder_stats(self):
305307
except RuntimeError:
306308
return # Stats not enabled
307309

310+
ANSIColors = get_colors()
311+
308312
print(f"\n{ANSIColors.BOLD_BLUE}{'='*50}{ANSIColors.RESET}")
309313
print(f"{ANSIColors.BOLD_BLUE}Unwinder Statistics:{ANSIColors.RESET}")
310314

@@ -399,6 +403,8 @@ def _print_binary_stats(self, collector):
399403
except (ValueError, RuntimeError):
400404
return # Collector closed or stats unavailable
401405

406+
ANSIColors = get_colors()
407+
402408
print(f" {ANSIColors.CYAN}Binary Encoding:{ANSIColors.RESET}")
403409

404410
repeat_records = stats.get('repeat_records', 0)

‎Lib/test/test_profiling/test_sampling_profiler/test_profiler.py‎

Lines changed: 11 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -17,7 +17,7 @@
1717
"Test only runs when _remote_debugging is available"
1818
)
1919

20-
from test.support import force_not_colorized_test_class
20+
from test.support import force_colorized, force_not_colorized_test_class
2121

2222

2323
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):
484484
self.assertIn("func1", result)
485485
self.assertIn("func2", result)
486486
self.assertIn("func3", result)
487+
self.assertNotIn("\x1b[", result)
488+
489+
@force_colorized
490+
def test_print_sampled_stats_colorized(self):
491+
with io.StringIO() as output, mock.patch("sys.stdout", output):
492+
print_sampled_stats(self.mock_stats)
493+
result = output.getvalue()
494+
495+
self.assertIn("\x1b[1;34mProfile Stats:", result)
487496

488497
def test_print_sampled_stats_sorting(self):
489498
"""Test different sorting options."""
@@ -751,6 +760,7 @@ def test_print_sampled_stats_sort_by_name(self):
751760
and not "calls" in line # Skip summary lines
752761
and not "total time" in line # Skip summary lines
753762
and not "cumulative time" in line
763+
and not "filename:lineno(function)" in line # Skip header line
754764
): # Skip summary lines
755765
data_lines.append(line)
756766

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,3 @@
1+
The :mod:`profiling.sampling` profiler now respects ``NO_COLOR``,
2+
``FORCE_COLOR``, and ``PYTHON_COLORS`` and avoids colored output when
3+
not writing to a terminal.

0 commit comments

Comments
 (0)