Add DesignGrid subclass for span-aware grid reductions

DesignGrid extends Grid to handle items spanning rectangular regions,
with span-aware reduce_cols, reduce_rows, and items_on_edge methods.
Items contribute their measurement divided by their span to each
affected row/column.

Author: has2k1

plotnine/_mpl/layout_manager/_grid.py

@@ -13,6 +13,9 @@
 
 T = TypeVar("T")
 
+Rect = tuple[int, int, int, int]
+"""Inclusive (r0, r1, c0, c1) rectangle in grid coordinates."""
+
 
 @dataclass
 class Grid(Generic[T]):
@@ -175,3 +178,101 @@ def items_on_edge(
         """
         cells = self[idx, :] if side in ("top", "bottom") else self[:, idx]
         return [n for n in cells if n is not None]
+
+
+class DesignGrid(Grid[T]):
+    """
+    Grid where items span rectangular regions
+
+    Each item is associated with an inclusive rectangle
+    `(r0, r1, c0, c1)` and placed at every cell within it; this
+    keeps base-class `__getitem__`, `iter_rows`, and `iter_cols`
+    working as in `Grid`. The reductions are overridden to be
+    span-aware: an item spanning multiple columns contributes its
+    measurement divided by its column span to each column it
+    covers (and analogously for rows).
+
+    Parameters
+    ----------
+    nrow
+        Number of rows in the grid.
+    ncol
+        Number of columns in the grid.
+    items
+        Items to place. One per rectangle, in the order rectangles
+        appear in `rects`.
+    rects
+        Inclusive `(r0, r1, c0, c1)` rectangle for each item.
+        Trusted: overlap and shape are not validated here.
+    """
+
+    # Bypass Grid's dataclass __init__ — rectangle expansion is a
+    # different placement scheme than row/col-major.
+    def __init__(
+        self,
+        nrow: int,
+        ncol: int,
+        items: Sequence[T],
+        rects: Sequence[Rect],
+    ):
+        if len(items) != len(rects):
+            raise ValueError(
+                f"Got {len(items)} items but {len(rects)} rectangles"
+            )
+        self._grid = np.empty((nrow, ncol), dtype=object)
+        self._items: list[T] = list(items)
+        self._rects: list[Rect] = list(rects)
+        # Place each item at every cell of its rectangle so the base
+        # class's __getitem__ / iter_rows / iter_cols keep working.
+        for item, (r0, r1, c0, c1) in zip(self._items, self._rects):
+            self._grid[r0 : r1 + 1, c0 : c1 + 1] = item
+
+    def reduce_cols(
+        self,
+        fn: Callable[[T], float],
+        default: float,
+    ) -> list[float]:
+        # An item spanning multiple columns shares its measurement
+        # across the columns it covers: fn(item) / colspan goes into
+        # each. Then per-column max as in Grid.reduce_cols.
+        out: list[float] = []
+        for c in range(self._grid.shape[1]):
+            contribs = [
+                fn(item) / (c1 - c0 + 1)
+                for item, (_, _, c0, c1) in zip(self._items, self._rects)
+                if c0 <= c <= c1
+            ]
+            out.append(max(contribs) if contribs else default)
+        return out
+
+    def reduce_rows(
+        self,
+        fn: Callable[[T], float],
+        default: float,
+    ) -> list[float]:
+        # Mirror of reduce_cols: fn(item) / rowspan into each row the
+        # item covers, then per-row max.
+        out: list[float] = []
+        for r in range(self._grid.shape[0]):
+            contribs = [
+                fn(item) / (r1 - r0 + 1)
+                for item, (r0, r1, _, _) in zip(self._items, self._rects)
+                if r0 <= r <= r1
+            ]
+            out.append(max(contribs) if contribs else default)
+        return out
+
+    def items_on_edge(
+        self,
+        side: Literal["top", "bottom", "left", "right"],
+        idx: int,
+    ) -> list[T]:
+        # An item's top/bottom edge is its r0/r1; left/right is c0/c1.
+        # Match the requested edge to idx exactly — not "the item is
+        # present at row/col idx", which would include spanned cells.
+        out: list[T] = []
+        for item, (r0, r1, c0, c1) in zip(self._items, self._rects):
+            edge = {"top": r0, "bottom": r1, "left": c0, "right": c1}[side]
+            if edge == idx:
+                out.append(item)
+        return out

tests/test_grid.py

@@ -1,6 +1,6 @@
 import pytest
 
-from plotnine._mpl.layout_manager._grid import Grid
+from plotnine._mpl.layout_manager._grid import DesignGrid, Grid
 
 
 def test_reduce_cols_basic():
@@ -107,3 +107,95 @@ def fn(n):
         return n
 
     grid.reduce_cols(fn, default=0)
+
+
+def test_design_grid_no_spans_matches_grid():
+    # Three single-cell items; each span is 1×1 so fn(item)/1 = fn(item).
+    # Mirrors what a plain Grid would do.
+    items = [1, 2, 3]
+    rects = [(0, 0, 0, 0), (0, 0, 1, 1), (1, 1, 0, 0)]
+    grid = DesignGrid[int](2, 2, items, rects)
+    assert grid.reduce_cols(lambda n: n, default=0) == [3, 2]
+    assert grid.reduce_rows(lambda n: n, default=0) == [2, 3]
+
+
+def test_design_grid_colspan_divides_contribution():
+    # Item spans both columns of a 1×2 grid.
+    grid = DesignGrid[int](1, 2, [10], [(0, 0, 0, 1)])
+    assert grid.reduce_cols(lambda n: n, default=0) == [5.0, 5.0]
+    assert grid.reduce_rows(lambda n: n, default=0) == [10.0]
+
+
+def test_design_grid_rowspan_divides_contribution():
+    # Item spans both rows of a 2×1 grid.
+    grid = DesignGrid[int](2, 1, [10], [(0, 1, 0, 0)])
+    assert grid.reduce_rows(lambda n: n, default=0) == [5.0, 5.0]
+    assert grid.reduce_cols(lambda n: n, default=0) == [10.0]
+
+
+def test_design_grid_square_span():
+    # Item spans the full 2×2 grid: fn / colspan = fn / rowspan = fn/2.
+    grid = DesignGrid[int](2, 2, [12], [(0, 1, 0, 1)])
+    assert grid.reduce_cols(lambda n: n, default=0) == [6.0, 6.0]
+    assert grid.reduce_rows(lambda n: n, default=0) == [6.0, 6.0]
+
+
+def test_design_grid_max_across_contributors():
+    # Col 0: contributions [10/1, 4/2] = [10, 2] → 10.
+    # Col 1: contributions [4/2]       = [2]     → 2.
+    items = [10, 4]
+    rects = [(0, 0, 0, 0), (1, 1, 0, 1)]
+    grid = DesignGrid[int](2, 2, items, rects)
+    assert grid.reduce_cols(lambda n: n, default=0) == [10.0, 2.0]
+
+
+def test_design_grid_empty_row_default():
+    # Items only in row 0 of a 3×2 grid; rows 1 and 2 take the default.
+    items = [1, 2]
+    rects = [(0, 0, 0, 0), (0, 0, 1, 1)]
+    grid = DesignGrid[int](3, 2, items, rects)
+    assert grid.reduce_rows(lambda n: n, default=99) == [2, 99, 99]
+
+
+def test_design_grid_empty_column_default():
+    # Items only in col 0 of a 2×3 grid; cols 1 and 2 take the default.
+    items = [1, 2]
+    rects = [(0, 0, 0, 0), (1, 1, 0, 0)]
+    grid = DesignGrid[int](2, 3, items, rects)
+    assert grid.reduce_cols(lambda n: n, default=99) == [2, 99, 99]
+
+
+def test_design_grid_items_on_edge_top_uses_r0():
+    # Spanning item: top edge at r0=0, not r1=1.
+    grid = DesignGrid[int](2, 1, [5], [(0, 1, 0, 0)])
+    assert grid.items_on_edge("top", 0) == [5]
+    assert grid.items_on_edge("top", 1) == []
+
+
+def test_design_grid_items_on_edge_bottom_uses_r1():
+    # Same item: bottom edge at r1=1, not r0=0.
+    grid = DesignGrid[int](2, 1, [5], [(0, 1, 0, 0)])
+    assert grid.items_on_edge("bottom", 1) == [5]
+    assert grid.items_on_edge("bottom", 0) == []
+
+
+def test_design_grid_items_on_edge_left_right():
+    # Col-spanning item: left edge at c0=0, right edge at c1=2.
+    grid = DesignGrid[int](1, 3, [7], [(0, 0, 0, 2)])
+    assert grid.items_on_edge("left", 0) == [7]
+    assert grid.items_on_edge("left", 1) == []
+    assert grid.items_on_edge("right", 2) == [7]
+    assert grid.items_on_edge("right", 0) == []
+
+
+def test_design_grid_mismatched_lengths_raises():
+    with pytest.raises(ValueError, match="2 items but 1 rectangles"):
+        DesignGrid[int](2, 2, [1, 2], [(0, 0, 0, 0)])
+
+
+def test_design_grid_indexing_returns_item_at_every_spanned_cell():
+    # Spanning item must appear at every cell of its rect so that
+    # base-class iter_rows / iter_cols continue to expose it.
+    grid = DesignGrid[int](2, 1, [5], [(0, 1, 0, 0)])
+    assert grid[0, 0] == 5
+    assert grid[1, 0] == 5