Merge pull request #26 from kdkavanagh/column-fit

Add auto column width mode

Author: kdkavanagh

src/dt_browser/browser.py

@@ -329,6 +329,7 @@ class DtBrowser(Widget):  # pylint: disable=too-many-public-methods,too-many-ins
         Binding("G", "last_row", "Jump to bottom", show=False),
         Binding("C", "show_colors", "Colors...", key_display="shift+C"),
         ("ctrl+s", "show_save", "Save dataframe as..."),
+        ("w", "toggle_auto_width", "Auto Width"),
     ]
 
     color_by: reactive[tuple[str, ...]] = reactive(tuple(), init=False)
@@ -571,6 +572,11 @@ def action_toggle_bookmark(self):
     async def action_toggle_row_detail(self):
         self.show_row_detail = not self.show_row_detail
 
+    def action_toggle_auto_width(self) -> None:
+        table = self.query_one("#main_table", CustomTable)
+        table.auto_width = not table.auto_width
+        self.notify("Auto column width: " + ("ON" if table.auto_width else "OFF"), timeout=2)
+
     async def action_last_row(self):
         table = self.query_one("#main_table", CustomTable)
         coord = table.cursor_coordinate

src/dt_browser/custom_table.py

@@ -192,6 +192,7 @@ def control(self) -> CustomTable:
     COMPONENT_CLASSES: ClassVar[set[str]] = {"datatable--header", "datatable--cursor", "datatable--even-row"}
 
     cursor_coordinate: Reactive[Coordinate] = Reactive(Coordinate(0, 0), repaint=False)
+    auto_width: Reactive[bool] = Reactive(False, repaint=False)
 
     def __init__(
         self,
@@ -221,6 +222,8 @@ def __init__(
 
         self._render_header_and_table: tuple[Strip, pl.DataFrame] | None = None
         self._dirty = True
+        self._full_widths: dict[str, int] = {}
+        self._auto_width_visible_range: tuple[int, int] | None = None
 
         self.set_dt(dt, metadata_dt)
 
@@ -268,6 +271,8 @@ def set_dt(self, dt: pl.DataFrame, metadata_dt: pl.DataFrame):
         self._dt = dt
         self._metadata_dt = metadata_dt
         self._set_widths({x: max(len(x), self._measure(self._dt[x])) for x in self._dt.columns})
+        self._full_widths = self._widths.copy()
+        self._auto_width_visible_range = None
         self._render_header_and_table = None
         self._formatters = {x: self._build_cast_expr(x, padding=self._widths[x]) for x in self._dt.columns}
         self._build_header_contents()
@@ -282,6 +287,24 @@ def _set_widths(self, widths: dict[str, int]):
             for k, v in zip(self._dt.columns, accumulate(x + COL_PADDING for x in self._widths.values()), strict=False)
         }
 
+    def _compute_auto_widths(self, scroll_y: int, dt_height: int) -> dict[str, int]:
+        """Compute column widths based only on currently visible rows."""
+        widths = {}
+        for col in self._dt.columns:
+            visible_slice = self._dt[col].slice(scroll_y, dt_height)
+            data_width = self._measure(visible_slice) if len(visible_slice) > 0 else 0
+            widths[col] = max(len(col), data_width)
+        return widths
+
+    def watch_auto_width(self, value: bool) -> None:
+        if not value:
+            self._set_widths(self._full_widths)
+            self._formatters = {x: self._build_cast_expr(x, padding=self._widths[x]) for x in self._dt.columns}
+            self._build_header_contents()
+        self._auto_width_visible_range = None
+        self._render_header_and_table = None
+        self.refresh(repaint=True)
+
     def render_line(self, y, *_):
         if y >= len(self._lines):
             pad = " " * (self.content_region.width)
@@ -357,6 +380,7 @@ def on_resize(self, event: events.Resize):
 
     def _ensure_cursor(self, allow_refresh: bool = True):
         self._render_header_and_table = None
+        self._auto_width_visible_range = None
 
         max_idx = self.cursor_coordinate.column
         while not self._is_col_visible(max_idx) and max_idx > 0:
@@ -496,10 +520,25 @@ def render_header_and_table(self):
             self._dirty = True
             scroll_x, scroll_y = self.scroll_offset
 
-            cols_to_render: list[str] = []
             effective_width = self.scrollable_content_region.width
             if effective_width <= 2:
                 return (Strip([]), pl.DataFrame())
+
+            dt_height = self.window_region.height - HEADER_HEIGHT
+
+            if self.auto_width:
+                visible_range = (scroll_y, dt_height)
+                if visible_range != self._auto_width_visible_range:
+                    self._auto_width_visible_range = visible_range
+                    new_widths = self._compute_auto_widths(scroll_y, dt_height)
+                    if new_widths != self._widths:
+                        self._set_widths(new_widths)
+                        self._formatters = {
+                            x: self._build_cast_expr(x, padding=self._widths[x]) for x in self._dt.columns
+                        }
+                        self._build_header_contents()
+
+            cols_to_render: list[str] = []
             truncate_last: int | None = None
             for x in self._dt.columns:
                 min_offset = self._cum_widths[x] - scroll_x
@@ -518,7 +557,6 @@ def render_header_and_table(self):
             if not cols_to_render:
                 return (Strip([]), pl.DataFrame())
 
-            dt_height = self.window_region.height - HEADER_HEIGHT
             base_header, header_width = self._build_base_header(cols_to_render)
             excess = self.scrollable_content_region.width - header_width
             header = Strip(base_header + (self._header_pad * (excess)))
@@ -584,6 +622,7 @@ def build_selector(cols: list[str], needed_padding: int = 0):
                 )
             else:
                 cursor_col_idx = self.cursor_coordinate.column - self._dt.columns.index(visible_cols[0])
+                cursor_col_idx = max(0, min(cursor_col_idx, len(visible_cols) - 1))
 
                 cols_before_selected: list[str] = visible_cols[0:cursor_col_idx]
                 sel_col = visible_cols[cursor_col_idx]

tests/__snapshots__/test_auto_width/test_snap_auto_width_off.svg

@@ -0,0 +1,238 @@
+import polars as pl
+from textual.pilot import Pilot
+
+from dt_browser.browser import DtBrowserApp
+from dt_browser.custom_table import CustomTable
+
+
+def _make_varying_width_app(num_rows: int = 50) -> DtBrowserApp:
+    """Create an app where row data varies in width so auto-width has a visible effect.
+
+    Early rows have short values, later rows have long values.
+    """
+    df = pl.DataFrame(
+        {
+            "short_col": [f"s{i}" for i in range(num_rows)],
+            "growing_col": [f"{'x' * (i + 1)}" for i in range(num_rows)],
+            "fixed_col": [f"fixed_{i:04d}" for i in range(num_rows)],
+        }
+    )
+    return DtBrowserApp("test", df)
+
+
+def _make_app_simple(num_rows: int = 20) -> DtBrowserApp:
+    df = pl.DataFrame(
+        {
+            "name": [f"item_{i}" for i in range(num_rows)],
+            "value": list(range(num_rows)),
+            "score": [round(i * 1.5, 1) for i in range(num_rows)],
+            "category": [f"cat_{i % 5}" for i in range(num_rows)],
+        }
+    )
+    return DtBrowserApp("test", df)
+
+
+async def test_auto_width_toggle():
+    """Toggling auto_width on and off updates the reactive property."""
+    app = _make_varying_width_app()
+    async with app.run_test(size=(120, 30)) as pilot:
+        await pilot.pause()
+        table = app.query_one("#main_table", CustomTable)
+
+        assert table.auto_width is False
+
+        await pilot.press("w")
+        await pilot.pause()
+        assert table.auto_width is True
+
+        await pilot.press("w")
+        await pilot.pause()
+        assert table.auto_width is False
+
+
+async def test_auto_width_narrows_columns():
+    """When auto_width is on, columns are narrower if visible rows have shorter data."""
+    app = _make_varying_width_app(num_rows=100)
+    async with app.run_test(size=(120, 30)) as pilot:
+        await pilot.pause()
+        table = app.query_one("#main_table", CustomTable)
+
+        # Full widths computed from all 100 rows (growing_col goes up to 100 chars)
+        full_widths = table._widths.copy()
+
+        await pilot.press("w")
+        await pilot.pause()
+
+        # Auto widths should be narrower for growing_col since visible rows are near the top
+        auto_widths = table._widths.copy()
+        assert auto_widths["growing_col"] < full_widths["growing_col"], (
+            f"Expected auto width ({auto_widths['growing_col']}) < full width ({full_widths['growing_col']}) "
+            f"for growing_col when viewing top rows"
+        )
+
+
+async def test_auto_width_updates_on_scroll():
+    """Scrolling to rows with wider data increases auto column widths."""
+    app = _make_varying_width_app(num_rows=100)
+    async with app.run_test(size=(120, 30)) as pilot:
+        await pilot.pause()
+        table = app.query_one("#main_table", CustomTable)
+
+        await pilot.press("w")
+        await pilot.pause()
+        widths_at_top = table._widths.copy()
+
+        # Scroll to the bottom where growing_col values are much wider
+        await pilot.press("G")
+        await pilot.pause()
+
+        widths_at_bottom = table._widths.copy()
+        assert widths_at_bottom["growing_col"] > widths_at_top["growing_col"], (
+            f"Expected wider growing_col at bottom ({widths_at_bottom['growing_col']}) "
+            f"than at top ({widths_at_top['growing_col']})"
+        )
+
+
+async def test_auto_width_restores_full_widths_on_toggle_off():
+    """Toggling auto_width off restores the original precomputed widths."""
+    app = _make_varying_width_app(num_rows=100)
+    async with app.run_test(size=(120, 30)) as pilot:
+        await pilot.pause()
+        table = app.query_one("#main_table", CustomTable)
+
+        full_widths = table._widths.copy()
+
+        # Toggle on
+        await pilot.press("w")
+        await pilot.pause()
+        assert table._widths != full_widths  # auto widths should differ
+
+        # Toggle off
+        await pilot.press("w")
+        await pilot.pause()
+        assert table._widths == full_widths, "Widths should be restored to full precomputed values"
+
+
+async def test_auto_width_respects_column_name_min_width():
+    """Auto width never makes a column narrower than its header name."""
+    app = _make_varying_width_app(num_rows=100)
+    async with app.run_test(size=(120, 30)) as pilot:
+        await pilot.pause()
+        table = app.query_one("#main_table", CustomTable)
+
+        await pilot.press("w")
+        await pilot.pause()
+
+        for col in table._dt.columns:
+            assert table._widths[col] >= len(col), (
+                f"Column '{col}' width ({table._widths[col]}) is less than header name length ({len(col)})"
+            )
+
+
+async def test_auto_width_skips_recompute_when_range_unchanged():
+    """Moving cursor within visible area does not trigger width recomputation."""
+    app = _make_varying_width_app(num_rows=100)
+    async with app.run_test(size=(120, 30)) as pilot:
+        await pilot.pause()
+        table = app.query_one("#main_table", CustomTable)
+
+        await pilot.press("w")
+        await pilot.pause()
+        widths_before = table._widths.copy()
+        range_before = table._auto_width_visible_range
+
+        # Move cursor within visible area
+        await pilot.press("down")
+        await pilot.pause()
+
+        assert table._auto_width_visible_range == range_before, "Visible range should not change for in-view cursor move"
+        assert table._widths == widths_before, "Widths should not change when cursor moves within visible area"
+
+
+async def test_auto_width_with_resize():
+    """Auto width recalculates when the terminal is resized."""
+    app = _make_varying_width_app(num_rows=100)
+    async with app.run_test(size=(120, 30)) as pilot:
+        await pilot.pause()
+        table = app.query_one("#main_table", CustomTable)
+
+        await pilot.press("w")
+        await pilot.pause()
+        range_before = table._auto_width_visible_range
+
+        # Resize changes dt_height, so auto widths should recompute
+        await pilot.resize_terminal(120, 20)
+        await pilot.pause()
+
+        # Force a render by accessing the property
+        _ = table.render_header_and_table
+
+        assert table._auto_width_visible_range != range_before, (
+            "Visible range should change after resize"
+        )
+
+
+async def test_auto_width_all_null_columns():
+    """Auto width does not crash when all visible columns have null values."""
+    df = pl.DataFrame(
+        {
+            "col_a": [None, None, None, None, None],
+            "col_b": [None, None, None, None, None],
+            "col_c": [None, None, None, None, None],
+        },
+        schema={"col_a": pl.Utf8, "col_b": pl.Utf8, "col_c": pl.Utf8},
+    )
+    app = DtBrowserApp("test", df)
+    async with app.run_test(size=(120, 30)) as pilot:
+        await pilot.pause()
+        table = app.query_one("#main_table", CustomTable)
+
+        # Switch to cell cursor mode then enable auto width
+        await pilot.press("w")
+        await pilot.pause()
+
+        # Should not crash — columns should be at least header-name width
+        for col in table._dt.columns:
+            assert table._widths[col] >= len(col)
+
+
+# --- Snapshot tests ---
+
+
+def test_snap_auto_width_off(snap_compare):
+    """Snapshot with auto width OFF (default)."""
+    assert snap_compare(_make_varying_width_app(num_rows=100), terminal_size=(120, 30))
+
+
+def test_snap_auto_width_on(snap_compare):
+    """Snapshot with auto width ON — columns should be narrower at top of data."""
+    assert snap_compare(
+        _make_varying_width_app(num_rows=100),
+        press=["w"],
+        terminal_size=(120, 30),
+    )
+
+
+def test_snap_auto_width_on_scrolled_bottom(snap_compare):
+    """Snapshot with auto width ON after scrolling to bottom — columns should be wider."""
+
+    async def run_before(pilot: Pilot) -> None:
+        await pilot.press("w")
+        await pilot.pause()
+        await pilot.press("G")
+        await pilot.pause()
+
+    assert snap_compare(
+        _make_varying_width_app(num_rows=100),
+        run_before=run_before,
+        terminal_size=(120, 30),
+    )
+
+
+def test_snap_auto_width_toggled_off_after_on(snap_compare):
+    """Snapshot after toggling auto width ON then OFF — should match original layout."""
+    assert snap_compare(
+        _make_varying_width_app(num_rows=100),
+        press=["w", "w"],
+        terminal_size=(120, 30),
+    )