fix: Respect header in last optimization step

CHANGELOG.md

@@ -9,6 +9,11 @@ The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.0.0/),
 ### Fixed
 
 - Fixed an issue where dynamic arrangement failed when setting the table to the exact width of the content [#90](https://github.com/Nukesor/comfy-table/issues/90).
+- The header size is now properly respected in the final optimization step [#90](https://github.com/Nukesor/comfy-table/issues/90).
+    Previously, this wasn't the case and lead to weird formatting behavior when both of the following were true
+  - Dynamic content adjustment was active.
+  - The table didn't fit into the the available space.
+  - The header of a row was longer than its content.
 
 ## [6.1.0] - 2022-08-28
 

src/table.rs

@@ -560,6 +560,37 @@ impl Table {
         }
     }
 
+    /// Get a mutable iterator over cells of a column, including the header cell.
+    /// The header cell will be the very first cell return, if a header is set.
+    /// The iterator returns a nested Option<Option<Cell>>, since there might be
+    /// rows that are missing this specific Cell.
+    ///
+    /// ```
+    /// use comfy_table::Table;
+    /// let mut table = Table::new();
+    /// table.set_header(&vec!["A", "B"]);
+    /// table.add_row(&vec!["First", "Second"]);
+    /// table.add_row(&vec!["Third"]);
+    /// table.add_row(&vec!["Fourth", "Fifth"]);
+    ///
+    /// // Create an iterator over the second column
+    /// let mut cell_iter = table.column_cells_with_header_iter(1);
+    /// assert_eq!(cell_iter.next().unwrap().unwrap().content(), "B");
+    /// assert_eq!(cell_iter.next().unwrap().unwrap().content(), "Second");
+    /// assert!(cell_iter.next().unwrap().is_none());
+    /// assert_eq!(cell_iter.next().unwrap().unwrap().content(), "Fifth");
+    /// assert!(cell_iter.next().is_none());
+    /// ```
+    pub fn column_cells_with_header_iter(&self, column_index: usize) -> ColumnCellsWithHeaderIter {
+        ColumnCellsWithHeaderIter {
+            header_checked: false,
+            header: &self.header,
+            rows: &self.rows,
+            column_index,
+            row_index: 0,
+        }
+    }
+
     /// Reference to a specific row
     pub fn row(&self, index: usize) -> Option<&Row> {
         self.rows.get(index)
@@ -679,12 +710,47 @@ impl<'a> Iterator for ColumnCellIter<'a> {
         if let Some(row) = self.rows.get(self.row_index) {
             self.row_index += 1;
 
-            // Check if the row has the requested column.
-            if let Some(cell) = row.cells.get(self.column_index) {
-                return Some(Some(cell));
-            }
+            // Return the cell (if it exists).
+            return Some(row.cells.get(self.column_index));
+        }
+
+        None
+    }
+}
+
+/// An iterator over cells of a specific column.
+/// A dedicated struct is necessary, as data is usually handled by rows and thereby stored in
+/// `Table::rows`. This type is returned by [Table::column_cells_iter].
+pub struct ColumnCellsWithHeaderIter<'a> {
+    header_checked: bool,
+    header: &'a Option<Row>,
+    rows: &'a [Row],
+    column_index: usize,
+    row_index: usize,
+}
+
+impl<'a> Iterator for ColumnCellsWithHeaderIter<'a> {
+    type Item = Option<&'a Cell>;
+    fn next(&mut self) -> Option<Option<&'a Cell>> {
+        // Get the header as the first cell
+        if !self.header_checked {
+            self.header_checked = true;
+
+            return match self.header {
+                Some(header) => {
+                    // Return the cell (if it exists).
+                    Some(header.cells.get(self.column_index))
+                }
+                None => None,
+            };
+        }
+
+        // Check if there's a next row
+        if let Some(row) = self.rows.get(self.row_index) {
+            self.row_index += 1;
 
-            return Some(None);
+            // Return the cell (if it exists).
+            return Some(row.cells.get(self.column_index));
         }
 
         None

src/utils/arrangement/dynamic.rs

@@ -40,7 +40,9 @@ pub fn arrange(
         available_content_width(table, infos, visible_columns, table_width);
 
     #[cfg(feature = "debug")]
-    println!("Table width: {table_width}, Start remaining width {remaining_width}");
+    println!(
+        "dynamic::arrange: Table width: {table_width}, Start remaining width {remaining_width}"
+    );
 
     // Step 2-4.
     // Find all columns that require less space than the average.
@@ -93,8 +95,8 @@ pub fn arrange(
 
     #[cfg(feature = "debug")]
     {
-        println!("After optimize: {infos:#?}",);
-        println!("Remaining width {remaining_width}, column {remaining_columns}",);
+        println!("dynamic::arrange: After optimize: {infos:#?}",);
+        println!("dynamic::arrange: Remaining width {remaining_width}, column {remaining_columns}",);
     }
 
     // Early exit and one branch of Part 7.
@@ -107,7 +109,7 @@ pub fn arrange(
         {
             use_full_width(infos, remaining_width);
             #[cfg(feature = "debug")]
-            println!("After full width: {infos:#?}");
+            println!("dynamic::arrange: After full width: {infos:#?}");
         }
         return;
     }
@@ -121,7 +123,7 @@ pub fn arrange(
     distribute_remaining_space(&table.columns, infos, remaining_width, remaining_columns);
 
     #[cfg(feature = "debug")]
-    println!("After distribute: {infos:#?}");
+    println!("dynamic::arrange: After distribute: {infos:#?}");
 }
 
 /// Step 1
@@ -249,8 +251,8 @@ fn find_columns_that_fit_into_average(
 
                     #[cfg(feature = "debug")]
                     println!(
-                        "dynamic::find_columns_that_fit_into_average: Fixed column {} via MaxWidth constraint with size {}",
-                        column.index, width
+                        "dynamic::find_columns_that_fit_into_average: Fixed column {} via MaxWidth constraint with size {}, as it's bigger than average {}",
+                        column.index, width, average_space
                     );
 
                     // Continue with new recalculated width
@@ -273,8 +275,8 @@ fn find_columns_that_fit_into_average(
 
                 #[cfg(feature = "debug")]
                 println!(
-                    "dynamic::find_columns_that_fit_into_average: Fixed column {} as it's smaller than average with size {}",
-                    column.index, max_column_width
+                    "dynamic::find_columns_that_fit_into_average: Fixed column {} with size {}, as it's smaller than average {}",
+                    column.index, max_column_width, average_space
                 );
 
                 // Continue with new recalculated width
@@ -294,10 +296,13 @@ fn find_columns_that_fit_into_average(
 
 /// Step 5
 ///
-/// Determine, whether there are any columns that have a higher LowerBoundary than the currently
-/// available average width.
+/// Determine, whether there are any columns that are allowed to occupy more width than the current
+/// `average_space` via a [LowerBoundary] constraint.
+///
+/// These columns will then get fixed to the width specified in the [LowerBoundary] constraint.
 ///
-/// These columns will then get fixed to the specified amount.
+/// I.e. if a column has to have at least 10 characters, but the average width left for a column is
+/// only 6, we fix the column to this 10 character minimum!
 fn enforce_lower_boundary_constraints(
     table: &Table,
     infos: &mut DisplayInfos,
@@ -337,6 +342,12 @@ fn enforce_lower_boundary_constraints(
         let info = ColumnDisplayInfo::new(column, width);
         infos.insert(column.index, info);
 
+        #[cfg(feature = "debug")]
+        println!(
+            "dynamic::enforce_lower_boundary_constraints: Fixed column {} to min constraint width {}",
+            column.index, width
+        );
+
         // Continue with new recalculated width
         remaining_width = remaining_width.saturating_sub(width.into());
         remaining_columns -= 1;
@@ -379,6 +390,12 @@ fn optimize_space_after_split(
     // Calculate the average space that remains for each column.
     let mut average_space = remaining_width / remaining_columns;
 
+    #[cfg(feature = "debug")]
+    println!(
+        "dynamic::optimize_space_after_split: Start with average_space {}",
+        average_space
+    );
+
     // Do this as long as we find a smaller column
     while found_smaller {
         found_smaller = false;
@@ -390,6 +407,12 @@ fn optimize_space_after_split(
 
             let longest_line = longest_line_after_split(average_space, column, table);
 
+            #[cfg(feature = "debug")]
+            println!(
+                "dynamic::optimize_space_after_split: Longest line after split for column {} is {}",
+                column.index, longest_line
+            );
+
             // If there's a considerable amount space left after splitting, we freeze the column and
             // set its content width to the calculated post-split width.
             let remaining_space = average_space.saturating_sub(longest_line);
@@ -404,6 +427,12 @@ fn optimize_space_after_split(
                     break;
                 }
                 average_space = remaining_width / remaining_columns;
+
+                #[cfg(feature = "debug")]
+                println!(
+                    "dynamic::optimize_space_after_split: average_space is now {}",
+                    average_space
+                );
                 found_smaller = true;
             }
         }
@@ -424,12 +453,11 @@ fn longest_line_after_split(average_space: usize, column: &Column, table: &Table
     let mut column_lines = Vec::new();
 
     // Iterate
-    for cell in table.column_cells_iter(column.index) {
+    for cell in table.column_cells_with_header_iter(column.index) {
         // Only look at rows that actually contain this cell.
-        let cell = if let Some(cell) = cell {
-            cell
-        } else {
-            continue;
+        let cell = match cell {
+            Some(cell) => cell,
+            None => continue,
         };
 
         let delimiter = delimiter(table, column, cell);
@@ -443,6 +471,13 @@ fn longest_line_after_split(average_space: usize, column: &Column, table: &Table
         for line in cell.content.iter() {
             if line.width() > average_space {
                 let mut splitted = split_line(line, &info, delimiter);
+
+                #[cfg(feature = "debug")]
+                println!(
+                    "dynamic::longest_line_after_split: Splitting line with width {}. Original:\n    {}\nSplitted:\n    {:?}",
+                    line.width(), line, splitted
+                );
+
                 column_lines.append(&mut splitted);
             } else {
                 column_lines.push(line.into());

tests/all/constraints_test.rs

@@ -3,6 +3,8 @@ use comfy_table::Width::*;
 use comfy_table::*;
 use pretty_assertions::assert_eq;
 
+use super::assert_table_line_width;
+
 fn get_constraint_table() -> Table {
     let mut table = Table::new();
     table
@@ -225,6 +227,7 @@ fn max_100_percentage() {
 | smol                                 |
 +--------------------------------------+";
     println!("{expected}");
+    assert_table_line_width(&table, 40);
     assert_eq!("\n".to_string() + &table.to_string(), expected);
 }
 

tests/all/content_arrangement_test.rs

@@ -3,13 +3,8 @@ use pretty_assertions::assert_eq;
 use comfy_table::ColumnConstraint::*;
 use comfy_table::Width::*;
 use comfy_table::{ContentArrangement, Row, Table};
-use unicode_width::UnicodeWidthStr;
 
-fn assert_table_lines(table: &Table, count: usize) {
-    for line in table.lines() {
-        assert_eq!(line.width(), count);
-    }
-}
+use super::assert_table_line_width;
 
 /// Test the robustnes of the dynamic table arangement.
 #[test]
@@ -67,7 +62,7 @@ fn simple_dynamic_table() {
 |        | stuff |      |
 +--------+-------+------+";
     println!("{expected}");
-    assert!(table.lines().all(|line| line.width() == 25));
+    assert_table_line_width(&table, 25);
     assert_eq!("\n".to_string() + &table.to_string(), expected);
 }
 
@@ -123,7 +118,7 @@ fn table_with_truncate() {
 | the middle ... | text   |       |
 +----------------+--------+-------+";
     println!("{expected}");
-    assert_table_lines(&table, 35);
+    assert_table_line_width(&table, 35);
     assert_eq!("\n".to_string() + &table.to_string(), expected);
 }
 
@@ -154,7 +149,7 @@ fn distribute_space_after_split() {
 +-----------------------------------------+-----------------------------+------+";
     println!("{expected}");
 
-    assert_table_lines(&table, 80);
+    assert_table_line_width(&table, 80);
     assert_eq!("\n".to_string() + &table.to_string(), expected);
 }
 
@@ -178,7 +173,7 @@ fn unused_space_after_split() {
 | anotherverylongtext |
 +---------------------+";
     println!("{expected}");
-    assert_table_lines(&table, 23);
+    assert_table_line_width(&table, 23);
     assert_eq!("\n".to_string() + &table.to_string(), expected);
 }
 
@@ -201,7 +196,7 @@ fn dynamic_full_width_after_split() {
 | anotherverylongtexttesttestaa                  |
 +------------------------------------------------+";
     println!("{expected}");
-    assert_table_lines(&table, 50);
+    assert_table_line_width(&table, 50);
     assert_eq!("\n".to_string() + &table.to_string(), expected);
 }
 
@@ -225,7 +220,7 @@ fn dynamic_full_width() {
 | This is a short line              | small                | smol              |
 +-----------------------------------+----------------------+-------------------+";
     println!("{expected}");
-    assert_table_lines(&table, 80);
+    assert_table_line_width(&table, 80);
     assert_eq!("\n".to_string() + &table.to_string(), expected);
 }
 
@@ -265,7 +260,45 @@ fn dynamic_exact_width() {
 │ 5   ┆ 6   ┆ 36.0      │
 └─────┴─────┴───────────┘";
         println!("{expected}");
-        assert_table_lines(&table, 25);
+        assert_table_line_width(&table, 25);
         assert_eq!("\n".to_string() + &table.to_string(), expected);
     }
 }
+
+/// Test that the formatting works as expected, if the table is slightly smaller than the max width
+/// of the table.
+#[rstest::rstest]
+fn dynamic_slightly_smaller() {
+    let header = vec!["a\n---\ni64", "b\n---\ni64", "b_squared\n---\nf64"];
+    let rows = vec![
+        vec!["1", "2", "4.0"],
+        vec!["3", "4", "16.0"],
+        vec!["5", "6", "36.0"],
+    ];
+
+    let mut table = Table::new();
+    let table = table
+        .load_preset(comfy_table::presets::UTF8_FULL)
+        .set_content_arrangement(ContentArrangement::Dynamic)
+        .set_width(24);
+
+    table.set_header(header.clone()).add_rows(rows.clone());
+
+    println!("{table}");
+    let expected = "
+┌─────┬─────┬──────────┐
+│ a   ┆ b   ┆ b_square │
+│ --- ┆ --- ┆ d        │
+│ i64 ┆ i64 ┆ ---      │
+│     ┆     ┆ f64      │
+╞═════╪═════╪══════════╡
+│ 1   ┆ 2   ┆ 4.0      │
+├╌╌╌╌╌┼╌╌╌╌╌┼╌╌╌╌╌╌╌╌╌╌┤
+│ 3   ┆ 4   ┆ 16.0     │
+├╌╌╌╌╌┼╌╌╌╌╌┼╌╌╌╌╌╌╌╌╌╌┤
+│ 5   ┆ 6   ┆ 36.0     │
+└─────┴─────┴──────────┘";
+    println!("{expected}");
+    assert_table_line_width(&table, 24);
+    assert_eq!("\n".to_string() + &table.to_string(), expected);
+}