feat: show progress during format conversion

austin/format/__init__.py

@@ -1,12 +1,18 @@
 from contextlib import contextmanager
+from contextlib import nullcontext
 from enum import Enum
+import os
 from pathlib import Path
+import sys
+import threading
 import typing as t
 
 from austin.events import AustinEventIterator
 
 MOJO_MAGIC = b"MOJ"
 
+ProgressCallback = t.Callable[[float], None]
+
 
 class Mode(Enum):
     """Austin profiling mode."""
@@ -28,21 +34,68 @@ def from_metadata(cls, mode: str) -> "Mode":
 
 
 @contextmanager
-def open_profile(path: t.Union[str, Path]) -> t.Iterator[AustinEventIterator]:
+def file_progress(
+    f: t.IO, on_progress: ProgressCallback, interval: float = 0.5
+) -> t.Iterator[None]:
+    """Report the fraction of a file that has been read.
+
+    The ``on_progress`` callback is called periodically from a background
+    thread, and once with ``1.0`` when the context exits normally. The file
+    position is read from the OS, so this works with any reader of the file.
+    """
+    size = os.fstat(f.fileno()).st_size
+
+    def report() -> None:
+        while not done.wait(interval):
+            on_progress(os.lseek(f.fileno(), 0, os.SEEK_CUR) / size if size else 1.0)
+
+    done = threading.Event()
+    reporter = threading.Thread(target=report, daemon=True)
+    reporter.start()
+    try:
+        yield
+    finally:
+        done.set()
+        reporter.join()
+
+    on_progress(1.0)
+
+
+def stderr_progress(label: str) -> t.Optional[ProgressCallback]:
+    """Get a progress callback that prints to stderr, if it is a terminal."""
+    if not sys.stderr.isatty():
+        return None
+
+    def on_progress(fraction: float) -> None:
+        line = f"{label} {fraction:.0%}"
+        if fraction >= 1.0:
+            # Done: clear the progress line
+            line = " " * len(line)
+        print(f"\r{line}\r", end="", file=sys.stderr, flush=True)
+
+    return on_progress
+
+
+@contextmanager
+def open_profile(
+    path: t.Union[str, Path], on_progress: t.Optional[ProgressCallback] = None
+) -> t.Iterator[AustinEventIterator]:
     """Open a MOJO or a collapsed stack profile file.
 
     The format is detected from the content of the file. The returned event
-    iterator is only valid within the context.
+    iterator is only valid within the context. If ``on_progress`` is given, it
+    is called with the fraction of the file read so far (see `file_progress`).
     """
     from austin.format.collapsed_stack import AustinFileReader
     from austin.format.mojo import MojoStreamReader
 
     with open(path, "rb") as f:
         is_mojo = f.read(len(MOJO_MAGIC)) == MOJO_MAGIC
 
-    if is_mojo:
-        with open(path, "rb") as f:
-            yield MojoStreamReader(f)
-    else:
-        with open(path) as f:
-            yield AustinFileReader(f)
+    with open(path, "rb" if is_mojo else "r") as f:
+        with file_progress(f, on_progress) if on_progress else nullcontext():
+            yield (
+                MojoStreamReader(f)
+                if is_mojo
+                else AustinFileReader(t.cast(t.TextIO, f))
+            )

austin/format/pprof/__main__.py

@@ -27,6 +27,7 @@
 from austin.events import AustinMetadata
 from austin.events import AustinSample
 from austin.format import open_profile
+from austin.format import stderr_progress
 from austin.format.pprof import PProf
 
 
@@ -56,7 +57,8 @@ def main() -> None:
     args = arg_parser.parse_args()
 
     try:
-        with open_profile(args.input) as fin:
+        progress = stderr_progress(f"Reading {args.input} ...")
+        with open_profile(args.input, progress) as fin:
             events = iter(fin)
             pprof = None
             for e in events:

austin/format/speedscope.py

@@ -38,6 +38,7 @@
 from austin.events import AustinSample
 from austin.format import Mode
 from austin.format import open_profile
+from austin.format import stderr_progress
 
 __version__ = "0.3.1"
 
@@ -223,7 +224,8 @@ def main() -> None:
     args = arg_parser.parse_args()
 
     try:
-        with open_profile(args.input) as fin:
+        progress = stderr_progress(f"Reading {args.input} ...")
+        with open_profile(args.input, progress) as fin:
             speedscope = None
             for event in fin:
                 if isinstance(event, AustinMetadata):

test/format/test_open_profile.py

@@ -0,0 +1,66 @@
+# This file is part of "austin-python" which is released under GPL.
+#
+# See file LICENCE or go to http://www.gnu.org/licenses/ for full license
+# details.
+#
+# austin-python is a Python wrapper around Austin, the CPython frame stack
+# sampler.
+#
+# Copyright (c) 2018-2026 Gabriele N. Tornetta <phoenix1987@gmail.com>.
+# All rights reserved.
+#
+# This program is free software: you can redistribute it and/or modify
+# it under the terms of the GNU General Public License as published by
+# the Free Software Foundation, either version 3 of the License, or
+# (at your option) any later version.
+#
+# This program is distributed in the hope that it will be useful,
+# but WITHOUT ANY WARRANTY; without even the implied warranty of
+# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
+# GNU General Public License for more details.
+# You should have received a copy of the GNU General Public License
+# along with this program.  If not, see <http://www.gnu.org/licenses/>.
+
+from pathlib import Path
+import time
+
+import pytest
+
+from austin.events import AustinSample
+from austin.format import file_progress
+from austin.format import open_profile
+from austin.format.collapsed_stack import AustinFileReader
+from austin.format.mojo import MojoStreamReader
+
+
+@pytest.mark.parametrize(
+    "name,reader", [("test.mojo", MojoStreamReader), ("test.austin", AustinFileReader)]
+)
+def test_open_profile_detects_format(datapath: Path, name, reader):
+    with open_profile(datapath / name) as events:
+        assert isinstance(events, reader)
+        assert any(isinstance(e, AustinSample) for e in events)
+
+
+@pytest.mark.parametrize("name", ["test.mojo", "test.austin"])
+def test_open_profile_progress(datapath: Path, name):
+    fractions: list[float] = []
+
+    with open_profile(datapath / name, fractions.append) as events:
+        for _ in events:
+            pass
+
+    assert fractions[-1] == 1.0
+    assert fractions == sorted(fractions)
+
+
+def test_file_progress_interval(tmp_path: Path):
+    path = tmp_path / "data"
+    path.write_bytes(b"x" * 100)
+    fractions: list[float] = []
+
+    with path.open("rb") as f, file_progress(f, fractions.append, interval=0.01):
+        f.raw.read(50)
+        time.sleep(0.1)
+
+    assert 0.5 in fractions and fractions[-1] == 1.0