From d7b98f4be6f943435032097177dd854e7a25e57c Mon Sep 17 00:00:00 2001 From: Kotakku Date: Thu, 26 Mar 2026 20:32:25 +0900 Subject: [PATCH] feat: add find_bags() and select_bag() TUI for interactive rosbag selection find_bags() discovers rosbag directories matching a glob pattern, sorted by the mtime of actual .db3/.mcap data files (not cache or directory timestamps). select_bag() provides a curses-based TUI for interactive selection with date display and keyboard navigation. Co-Authored-By: Claude Opus 4.6 (1M context) --- example/plot_twist_xyz.py | 6 +-- example/select_and_plot.py | 25 +++++++++ src/baglab/__init__.py | 5 +- src/baglab/io/__init__.py | 3 +- src/baglab/io/bag.py | 36 +++++++++++++ src/baglab/tui.py | 107 +++++++++++++++++++++++++++++++++++++ test/test_io.py | 45 ++++++++++++++++ 7 files changed, 222 insertions(+), 5 deletions(-) create mode 100644 example/select_and_plot.py create mode 100644 src/baglab/tui.py diff --git a/example/plot_twist_xyz.py b/example/plot_twist_xyz.py index 78c2c86..fa93de5 100644 --- a/example/plot_twist_xyz.py +++ b/example/plot_twist_xyz.py @@ -4,12 +4,12 @@ import matplotlib.pyplot as plt -import baglab +import baglab as bl -bag = baglab.load(Path(__file__).parent / "test_bag") +bag = bl.load(Path(__file__).parent / "test_bag") twist_df = bag["/test/twist"] -t = baglab.stamp_to_sec(twist_df, relative=True) +t = bl.stamp_to_sec(twist_df, relative=True) vel = twist_df.msg.twist.linear.df # columns: [x, y, z] fig, axes = plt.subplots(3, 1, sharex=True) diff --git a/example/select_and_plot.py b/example/select_and_plot.py new file mode 100644 index 0000000..a329eb4 --- /dev/null +++ b/example/select_and_plot.py @@ -0,0 +1,25 @@ +"""Select a rosbag via TUI and plot twist linear x, y, z over time.""" + +from pathlib import Path + +import matplotlib.pyplot as plt + +import baglab as bl + +bag_path = bl.select_bag(str(Path(__file__).parent / "*")) +print(f"Selected: {bag_path}") + +bag = bl.load(bag_path) +twist_df = bag["/test/twist"] + +t = bl.stamp_to_sec(twist_df, relative=True) +vel = twist_df.msg.twist.linear.df # columns: [x, y, z] + +fig, axes = plt.subplots(3, 1, sharex=True) +for ax, axis_name in zip(axes, ["x", "y", "z"]): + ax.plot(t, vel[axis_name]) + ax.set_ylabel(axis_name) +axes[-1].set_xlabel("time [s]") +fig.suptitle(f"twist.linear — {bag_path.name}") +plt.tight_layout() +plt.show() diff --git a/src/baglab/__init__.py b/src/baglab/__init__.py index 5088f54..2d49e15 100644 --- a/src/baglab/__init__.py +++ b/src/baglab/__init__.py @@ -37,9 +37,10 @@ from baglab.plot import plot_error_band, plot_step_response, plot_timeseries, plot_xy_trajectory from baglab.analysis import delay_estimate, stepinfo, tracking_error from baglab.diagnostics import latency_chain, message_gaps, topic_delay, topic_rate -from baglab.io import MsgAccessor, FieldGroup, align_origin, clear_cache, explode_array, find_time, has_mcap_backend, load, recv_time_to_sec, reindex_by_stamp, stamp_to_sec, time_slice +from baglab.io import MsgAccessor, FieldGroup, align_origin, clear_cache, explode_array, find_bags, find_time, has_mcap_backend, load, recv_time_to_sec, reindex_by_stamp, stamp_to_sec, time_slice from baglab.signal import diff, fft, integrate, lowpass, moving_average from baglab.stats import describe, rms +from baglab.tui import select_bag from baglab.geometry import ( align_time, angle_diff, @@ -73,6 +74,7 @@ "align_origin", "clear_cache", "explode_array", + "find_bags", "find_time", "has_mcap_backend", "load", @@ -118,6 +120,7 @@ "message_gaps", "topic_delay", "topic_rate", + "select_bag", "plot_error_band", "plot_step_response", "plot_timeseries", diff --git a/src/baglab/io/__init__.py b/src/baglab/io/__init__.py index 9e80b30..b2d3ecf 100644 --- a/src/baglab/io/__init__.py +++ b/src/baglab/io/__init__.py @@ -1,7 +1,7 @@ """I/O module for rosbag loading, field access, and timestamp utilities.""" from baglab.io.accessor import FieldGroup, MsgAccessor, explode_array -from baglab.io.bag import Bag, clear_cache, has_mcap_backend, load +from baglab.io.bag import Bag, clear_cache, find_bags, has_mcap_backend, load from baglab.io.stamp import align_origin, find_time, recv_time_to_sec, reindex_by_stamp, stamp_to_sec, time_slice from baglab.io.typesys import register_msg_types @@ -10,6 +10,7 @@ "align_origin", "clear_cache", "explode_array", + "find_bags", "FieldGroup", "MsgAccessor", "find_time", diff --git a/src/baglab/io/bag.py b/src/baglab/io/bag.py index aff38a6..9444076 100644 --- a/src/baglab/io/bag.py +++ b/src/baglab/io/bag.py @@ -165,6 +165,42 @@ def _bag_data_files(bag_path: Path) -> list[Path]: ) +def _bag_data_mtime(bag_path: Path) -> float: + """Return the newest mtime among data files (.db3/.mcap) in a bag directory.""" + data_files = _bag_data_files(bag_path) + if data_files: + return max(f.stat().st_mtime for f in data_files) + return bag_path.stat().st_mtime + + +def find_bags(pattern: str) -> list[Path]: + """Return bag paths matching a glob pattern, sorted by data-file modification time. + + Sorts by the modification time of the actual ``.db3`` or ``.mcap`` files + inside each bag directory (oldest first), ignoring cache or metadata + timestamps. + + Parameters + ---------- + pattern : str + Glob pattern (e.g. ``"/path/to/log_dir/*"``). + + Returns + ------- + list[Path] + Matched paths sorted ascending by data-file mtime. + Use ``[-1]`` for the latest bag. + """ + from glob import glob as _glob + + paths = [ + Path(p) for p in _glob(pattern) + if Path(p).is_dir() and _bag_data_files(Path(p)) + ] + paths.sort(key=_bag_data_mtime) + return paths + + def _compute_fingerprint(bag_path: Path) -> dict[str, dict]: """Compute fingerprint of bag data files for cache invalidation.""" result = {} diff --git a/src/baglab/tui.py b/src/baglab/tui.py new file mode 100644 index 0000000..7eb36d4 --- /dev/null +++ b/src/baglab/tui.py @@ -0,0 +1,107 @@ +"""Terminal UI utilities for interactive rosbag selection.""" + +from __future__ import annotations + +import curses +import datetime +from pathlib import Path + +from baglab.io.bag import _bag_data_mtime, find_bags + + +def select_bag(pattern: str) -> Path: + """Show a curses TUI to interactively select a rosbag directory. + + Bags matching *pattern* are listed with their data-file modification + timestamp, sorted oldest-first (newest at the bottom). Use arrow keys + or ``j``/``k`` to move, Enter to confirm, ``q`` to cancel. + + Parameters + ---------- + pattern : str + Glob pattern forwarded to :func:`baglab.find_bags`. + + Returns + ------- + Path + The selected bag directory path. + + Raises + ------ + FileNotFoundError + If no bags match *pattern*. + KeyboardInterrupt + If the user cancels with ``q`` or Ctrl-C. + """ + bags = find_bags(pattern) + if not bags: + raise FileNotFoundError(f"No bags found matching: {pattern}") + + # Pre-build display lines: " YYYY-MM-DD HH:MM:SS bag_name" + lines: list[tuple[str, Path]] = [] + for bag in bags: + mtime = _bag_data_mtime(bag) + dt = datetime.datetime.fromtimestamp(mtime) + ts = dt.strftime("%Y-%m-%d %H:%M:%S") + lines.append((ts, bag)) + + selected = len(lines) - 1 # default to newest (last) + + def _run(stdscr: curses.window) -> Path: + nonlocal selected + curses.curs_set(0) + curses.use_default_colors() + curses.init_pair(1, curses.COLOR_BLACK, curses.COLOR_CYAN) + + while True: + stdscr.erase() + max_y, max_x = stdscr.getmaxyx() + + title = f"Select a rosbag ({len(lines)} found) [↑↓/jk: move, Enter: select, q: cancel]" + stdscr.addnstr(0, 0, title, max_x - 1, curses.A_BOLD) + + # Scrollable area + list_start = 2 + visible = max_y - list_start + if visible <= 0: + visible = 1 + + # Keep selected item visible + if selected < 0: + selected = 0 + if selected >= len(lines): + selected = len(lines) - 1 + + # Calculate scroll offset + offset = 0 + if selected >= visible: + offset = selected - visible + 1 + + for i in range(offset, min(offset + visible, len(lines))): + ts, bag = lines[i] + row = list_start + (i - offset) + label = f" {ts} {bag.name}" + if i == selected: + label = f"> {ts} {bag.name}" + attr = curses.color_pair(1) | curses.A_BOLD + else: + attr = curses.A_NORMAL + stdscr.addnstr(row, 0, label, max_x - 1, attr) + + stdscr.refresh() + + key = stdscr.getch() + if key in (curses.KEY_UP, ord("k")): + selected = max(0, selected - 1) + elif key in (curses.KEY_DOWN, ord("j")): + selected = min(len(lines) - 1, selected + 1) + elif key in (curses.KEY_HOME, ord("g")): + selected = 0 + elif key in (curses.KEY_END, ord("G")): + selected = len(lines) - 1 + elif key in (curses.KEY_ENTER, 10, 13): + return lines[selected][1] + elif key in (ord("q"), 27): # q or Escape + raise KeyboardInterrupt("Selection cancelled") + + return curses.wrapper(_run) diff --git a/test/test_io.py b/test/test_io.py index 544b99d..b3ff3b9 100644 --- a/test/test_io.py +++ b/test/test_io.py @@ -1,5 +1,8 @@ """Tests for baglab.io module.""" +import os +import time + import pandas as pd import pytest @@ -7,6 +10,48 @@ from baglab.io import Bag, clear_cache +class TestFindBags: + """Tests for baglab.find_bags().""" + + def test_finds_test_bag(self, test_bag_path): + pattern = str(test_bag_path.parent / "*") + results = baglab.find_bags(pattern) + assert any(r == test_bag_path for r in results) + + def test_returns_list_of_paths(self, test_bag_path): + from pathlib import Path + + pattern = str(test_bag_path.parent / "*") + results = baglab.find_bags(pattern) + assert isinstance(results, list) + assert all(isinstance(p, Path) for p in results) + + def test_sorted_by_data_file_mtime(self, test_bag_path, tmp_path): + """Bags should be sorted by data-file mtime, not directory mtime.""" + import shutil + + # Create two bag copies with controlled data-file mtimes + bag_old = tmp_path / "bag_old" + bag_new = tmp_path / "bag_new" + shutil.copytree(test_bag_path, bag_old) + time.sleep(0.05) + shutil.copytree(test_bag_path, bag_new) + + # Touch data files in bag_old to make them newer + for f in bag_old.iterdir(): + if f.suffix in (".db3", ".mcap"): + f.touch() + + results = baglab.find_bags(str(tmp_path / "*")) + bag_names = [r.name for r in results] + # bag_new was copied later but bag_old has newer data files + assert bag_names.index("bag_new") < bag_names.index("bag_old") + + def test_empty_pattern(self, tmp_path): + results = baglab.find_bags(str(tmp_path / "nonexistent_*")) + assert results == [] + + class TestLoad: """Tests for baglab.load()."""