From 9ddf7abeeeb199e8e9f3189b75e2e0e2c037cebd Mon Sep 17 00:00:00 2001
From: Andrew Scott <andrewcscott007@gmail.com>
Date: Thu, 30 Mar 2023 03:01:22 -0500
Subject: [PATCH] feat: simple grid row color picker (#2519)

* feat: simple grid row color picker

* feat: allow painting initial grid row

* avoid coloring callback for striped when not needed

Co-authored-by: Emil Ernerfeldt <emil.ernerfeldt@gmail.com>

* fix: switch from rgba to color32 for row grid color assignment

* fix warning

* make `paint_row` private

---------

Co-authored-by: Emil Ernerfeldt <emil.ernerfeldt@gmail.com>
---
 crates/egui/src/grid.rs | 75 +++++++++++++++++++++++++++++------------
 1 file changed, 53 insertions(+), 22 deletions(-)

diff --git a/crates/egui/src/grid.rs b/crates/egui/src/grid.rs
index e84b6b42..c3427a96 100644
--- a/crates/egui/src/grid.rs
+++ b/crates/egui/src/grid.rs
@@ -46,6 +46,9 @@ impl State {
 
 // ----------------------------------------------------------------------------
 
+// type alias for boxed function to determine row color during grid generation
+type ColorPickerFn = Box<dyn Fn(usize, &Style) -> Option<Color32>>;
+
 pub(crate) struct GridLayout {
     ctx: Context,
     style: std::sync::Arc<Style>,
@@ -66,7 +69,7 @@ pub(crate) struct GridLayout {
     spacing: Vec2,
     min_cell_size: Vec2,
     max_cell_size: Vec2,
-    striped: bool,
+    color_picker: Option<ColorPickerFn>,
 
     // Cursor:
     col: usize,
@@ -101,7 +104,7 @@ impl GridLayout {
             spacing: ui.spacing().item_spacing,
             min_cell_size: ui.spacing().interact_size,
             max_cell_size: Vec2::INFINITY,
-            striped: false,
+            color_picker: None,
 
             col: 0,
             row: 0,
@@ -214,6 +217,20 @@ impl GridLayout {
         self.col += 1;
     }
 
+    fn paint_row(&mut self, cursor: &mut Rect, painter: &Painter) {
+        // handle row color painting based on color-picker function
+        let Some(color_picker) = self.color_picker.as_ref() else { return };
+        let Some(row_color)    = color_picker(self.row, &self.style) else { return };
+        let Some(height)       = self.prev_state.row_height(self.row) else {return };
+        // Paint background for coming row:
+        let size = Vec2::new(self.prev_state.full_width(self.spacing.x), height);
+        let rect = Rect::from_min_size(cursor.min, size);
+        let rect = rect.expand2(0.5 * self.spacing.y * Vec2::Y);
+        let rect = rect.expand2(2.0 * Vec2::X); // HACK: just looks better with some spacing on the sides
+
+        painter.rect_filled(rect, 2.0, row_color);
+    }
+
     pub(crate) fn end_row(&mut self, cursor: &mut Rect, painter: &Painter) {
         cursor.min.x = self.initial_available.min.x;
         cursor.min.y += self.spacing.y;
@@ -225,17 +242,7 @@ impl GridLayout {
         self.col = 0;
         self.row += 1;
 
-        if self.striped && self.row % 2 == 1 {
-            if let Some(height) = self.prev_state.row_height(self.row) {
-                // Paint background for coming row:
-                let size = Vec2::new(self.prev_state.full_width(self.spacing.x), height);
-                let rect = Rect::from_min_size(cursor.min, size);
-                let rect = rect.expand2(0.5 * self.spacing.y * Vec2::Y);
-                let rect = rect.expand2(2.0 * Vec2::X); // HACK: just looks better with some spacing on the sides
-
-                painter.rect_filled(rect, 2.0, self.style.visuals.faint_bg_color);
-            }
-        }
+        self.paint_row(cursor, painter);
     }
 
     pub(crate) fn save(&self) {
@@ -278,12 +285,12 @@ impl GridLayout {
 pub struct Grid {
     id_source: Id,
     num_columns: Option<usize>,
-    striped: Option<bool>,
     min_col_width: Option<f32>,
     min_row_height: Option<f32>,
     max_cell_size: Vec2,
     spacing: Option<Vec2>,
     start_row: usize,
+    color_picker: Option<ColorPickerFn>,
 }
 
 impl Grid {
@@ -292,15 +299,24 @@ impl Grid {
         Self {
             id_source: Id::new(id_source),
             num_columns: None,
-            striped: None,
             min_col_width: None,
             min_row_height: None,
             max_cell_size: Vec2::INFINITY,
             spacing: None,
             start_row: 0,
+            color_picker: None,
         }
     }
 
+    /// Setting this will allow for dynamic coloring of rows of the grid object
+    pub fn with_row_color<F>(mut self, color_picker: F) -> Self
+    where
+        F: Fn(usize, &Style) -> Option<Color32> + 'static,
+    {
+        self.color_picker = Some(Box::new(color_picker));
+        self
+    }
+
     /// Setting this will allow the last column to expand to take up the rest of the space of the parent [`Ui`].
     pub fn num_columns(mut self, num_columns: usize) -> Self {
         self.num_columns = Some(num_columns);
@@ -311,9 +327,17 @@ impl Grid {
     ///
     /// This can make a table easier to read.
     /// Default is whatever is in [`crate::Visuals::striped`].
-    pub fn striped(mut self, striped: bool) -> Self {
-        self.striped = Some(striped);
-        self
+    pub fn striped(self, striped: bool) -> Self {
+        if striped {
+            self.with_row_color(move |row, style| {
+                if row % 2 == 1 {
+                    return Some(style.visuals.faint_bg_color);
+                }
+                None
+            })
+        } else {
+            self
+        }
     }
 
     /// Set minimum width of each column.
@@ -364,14 +388,13 @@ impl Grid {
         let Self {
             id_source,
             num_columns,
-            striped,
             min_col_width,
             min_row_height,
             max_cell_size,
             spacing,
             start_row,
+            color_picker,
         } = self;
-        let striped = striped.unwrap_or(ui.visuals().striped);
         let min_col_width = min_col_width.unwrap_or_else(|| ui.spacing().interact_size.x);
         let min_row_height = min_row_height.unwrap_or_else(|| ui.spacing().interact_size.y);
         let spacing = spacing.unwrap_or_else(|| ui.spacing().item_spacing);
@@ -387,9 +410,10 @@ impl Grid {
         ui.allocate_ui_at_rect(max_rect, |ui| {
             ui.set_visible(prev_state.is_some()); // Avoid visible first-frame jitter
             ui.horizontal(|ui| {
-                let grid = GridLayout {
+                let is_color = color_picker.is_some();
+                let mut grid = GridLayout {
                     num_columns,
-                    striped,
+                    color_picker,
                     min_cell_size: vec2(min_col_width, min_row_height),
                     max_cell_size,
                     spacing,
@@ -397,6 +421,13 @@ impl Grid {
                     ..GridLayout::new(ui, id, prev_state)
                 };
 
+                // paint first incoming row
+                if is_color {
+                    let mut cursor = ui.cursor();
+                    let painter = ui.painter();
+                    grid.paint_row(&mut cursor, painter);
+                }
+
                 ui.set_grid(grid);
                 let r = add_contents(ui);
                 ui.save_grid();