diff --git a/flameshow/main.py b/flameshow/main.py index ea31f07..66d9426 100644 --- a/flameshow/main.py +++ b/flameshow/main.py @@ -2,6 +2,8 @@ import os import time +import sys + import click from flameshow import __version__ @@ -36,6 +38,15 @@ def setup_log(enabled, level, loglocation): } +def ensure_tty(): + if os.isatty(0): + return + + logger.info("stdin is not a tty, replace it to fd=2") + sys.stdin.close() + sys.stdin = os.fdopen(2) + + def run_app(verbose, log_to, profile_f, _debug_exit_after_rednder): log_level = LOG_LEVEL[verbose] setup_log(log_to is not None, log_level, log_to) @@ -48,6 +59,7 @@ def run_app(verbose, log_to, profile_f, _debug_exit_after_rednder): t01 = time.time() logger.info("Parse profile, took %.3fs", t01 - t0) + ensure_tty() app = FlameshowApp( profile, _debug_exit_after_rednder, diff --git a/flameshow/render/flamegraph.py b/flameshow/render/flamegraph.py index 5d326b0..7afa195 100644 --- a/flameshow/render/flamegraph.py +++ b/flameshow/render/flamegraph.py @@ -116,7 +116,6 @@ def generate_frame_maps(self, width, focused_stack_id): logger.info("frame maps: %s", frame_maps) def _generate_for_children(frame): - logger.debug("generate frame_maps for %s", frame) # generate for children my_maps = frame_maps[frame._id] for sample_i, my_map in enumerate(my_maps): @@ -248,13 +247,6 @@ def render_line(self, y: int) -> Strip: ) cursor += my_width - logger.debug( - "%s in line %d, frame_map=%s", - frame, - y, - frame_map, - ) - strip = Strip(segments) return strip diff --git a/tests/test_cli_click.py b/tests/test_cli_click.py index cadafd3..c4bc6f9 100644 --- a/tests/test_cli_click.py +++ b/tests/test_cli_click.py @@ -1,6 +1,6 @@ from click.testing import CliRunner -from flameshow.main import main -from unittest.mock import patch +from flameshow.main import main, ensure_tty +from unittest.mock import patch, MagicMock def test_print_version(): @@ -11,10 +11,29 @@ def test_print_version(): assert result.output == "1.2.3\n" -def test_run_app(): +@patch("flameshow.main.os") +def test_run_app(mock_os, data_dir): + mock_os.isatty.return_value = True runner = CliRunner() result = runner.invoke( - main, ["tests/pprof_data/profile-10seconds.out"], input="q" + main, [str(data_dir / "profile-10seconds.out")], input="q" ) assert result.exit_code == 0 + + +@patch("flameshow.main.sys") +@patch("flameshow.main.os") +def test_ensure_tty_when_its_not(mock_os, mock_sys): + mock_os.isatty.return_value = False + opened_fd = object() + mock_os.fdopen.return_value = opened_fd + + fake_stdin = MagicMock() + mock_sys.stdin = fake_stdin + + ensure_tty() + + fake_stdin.close.assert_called() + assert hasattr(mock_sys, "stdin") + assert mock_sys.stdin is opened_fd