Moved empty line check in table creator

2 files changed, 13 insertions, 4 deletions

diff --git a/cmd2/table_creator.py b/cmd2/table_creator.py
index d159ddbf..3b0db7b3 100644
--- a/cmd2/table_creator.py
+++ b/cmd2/table_creator.py
@@ -309,6 +309,11 @@ class TableCreator:
             if data_line_index > 0:
                 wrapped_buf.write('\n')
 
+            # If the last line is empty, then add a newline and stop
+            if data_line_index == len(data_str_lines) - 1 and not data_line:
+                wrapped_buf.write('\n')
+                break
+
             # Locate the styles in this line
             styles = utils.get_styles_in_text(data_line)
 
@@ -351,10 +356,6 @@ class TableCreator:
                 last_word = data_line_index == len(data_str_lines) - 1 and char_index == len(data_line)
                 add_word(cur_word_buf.getvalue(), last_word)
 
-            # If the last line is empty, then add a newline
-            elif data_line_index == len(data_str_lines) - 1:
-                wrapped_buf.write('\n')
-
             # Stop line loop if we've written to max_lines
             if total_lines == max_lines:
                 # If this isn't the last data line and there is space
diff --git a/tests/test_table_creator.py b/tests/test_table_creator.py
index 396016eb..ee15bb22 100644
--- a/tests/test_table_creator.py
+++ b/tests/test_table_creator.py
@@ -133,6 +133,14 @@ def test_blank_last_line():
     assert row == ('my line   \n'
                    '          ')
 
+    row_data = ['\n']
+    row = tc.generate_row(row_data=row_data)
+    assert row == '          '
+
+    row_data = ['']
+    row = tc.generate_row(row_data=row_data)
+    assert row == '          '
+
 
 def test_wrap_text():
     column_1 = Column("Col 1", width=10)