Added clearer exception handling to BorderedTable and SimpleTable

3 files changed, 36 insertions, 2 deletions

diff --git a/CHANGELOG.md b/CHANGELOG.md
index 0f231fec..b5df914d 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -1,3 +1,7 @@
+## 2.3.3 (TBD, 2021)
+* Enhancements
+  * Added clearer exception handling to `BorderedTable` and `SimpleTable`.
+
 ## 2.3.2 (November 22, 2021)
 * Bug Fixes
   * Fixed issue where a `ns_provider` could be passed `None` instead of its correct `cmd2.Cmd` or `CommandSet` value.
diff --git a/cmd2/table_creator.py b/cmd2/table_creator.py
index 6db07517..bb8504ba 100644
--- a/cmd2/table_creator.py
+++ b/cmd2/table_creator.py
@@ -440,7 +440,7 @@ class TableCreator:
         :param post_line: string to print after each line of a row. This can be used for padding after
                           the last cell's text and a right row border. (Defaults to blank)
         :return: row string
-        :raises: ValueError if data isn't the same length as self.cols
+        :raises: ValueError if row_data isn't the same length as self.cols
         :raises: TypeError if fill_char is more than one character (not including ANSI style sequences)
         :raises: ValueError if fill_char, pre_line, inter_cell, or post_line contains an unprintable
                  character like a newline
@@ -682,7 +682,11 @@ class SimpleTable(TableCreator):
 
         :param row_data: data with an entry for each column in the row
         :return: data row string
+        :raises: ValueError if row_data isn't the same length as self.cols
         """
+        if len(row_data) != len(self.cols):
+            raise ValueError("Length of row_data must match length of cols")
+
         fill_char = self.apply_data_bg(SPACE)
         inter_cell = self.apply_data_bg(self.column_spacing * SPACE)
 
@@ -969,7 +973,11 @@ class BorderedTable(TableCreator):
 
         :param row_data: data with an entry for each column in the row
         :return: data row string
+        :raises: ValueError if row_data isn't the same length as self.cols
         """
+        if len(row_data) != len(self.cols):
+            raise ValueError("Length of row_data must match length of cols")
+
         fill_char = self.apply_data_bg(SPACE)
 
         pre_line = self.apply_border_color('║') + self.apply_data_bg(self.padding * SPACE)
diff --git a/tests/test_table_creator.py b/tests/test_table_creator.py
index 52c4da8b..585ed62a 100644
--- a/tests/test_table_creator.py
+++ b/tests/test_table_creator.py
@@ -316,7 +316,7 @@ def test_generate_row_exceptions():
             tc.generate_row(row_data=row_data, is_header=False, **kwargs)
         assert "{} contains an unprintable character".format(arg) in str(excinfo.value)
 
-    # data with too many columns
+    # Data with too many columns
     row_data = ['Data 1', 'Extra Column']
     with pytest.raises(ValueError) as excinfo:
         tc.generate_row(row_data=row_data, is_header=False)
@@ -504,6 +504,17 @@ def test_simple_table_width():
     assert st.total_width() == 34
 
 
+def test_simple_generate_data_row_exceptions():
+    column_1 = Column("Col 1")
+    tc = SimpleTable([column_1])
+
+    # Data with too many columns
+    row_data = ['Data 1', 'Extra Column']
+    with pytest.raises(ValueError) as excinfo:
+        tc.generate_data_row(row_data=row_data)
+    assert "Length of row_data must match length of cols" in str(excinfo.value)
+
+
 def test_bordered_table_creation():
     column_1 = Column("Col 1", width=15)
     column_2 = Column("Col 2", width=15)
@@ -635,6 +646,17 @@ def test_bordered_table_width():
     assert bt.total_width() == 37
 
 
+def test_bordered_generate_data_row_exceptions():
+    column_1 = Column("Col 1")
+    tc = BorderedTable([column_1])
+
+    # Data with too many columns
+    row_data = ['Data 1', 'Extra Column']
+    with pytest.raises(ValueError) as excinfo:
+        tc.generate_data_row(row_data=row_data)
+    assert "Length of row_data must match length of cols" in str(excinfo.value)
+
+
 def test_alternating_table_creation():
     column_1 = Column("Col 1", width=15)
     column_2 = Column("Col 2", width=15)