Merge pull request #35 from laixintao/reopen-stdin

support passing profile data though PIPE, by reopen fd=2 if stdin is not a tty.

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,

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
 

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