DataTable issues with adding new columns to a table that has existing data in it

[mrllama123]

Hello, I'm working on a project that displays DynamoDB table (a nosql db) data called dyno-viewer

I encountered and issue where with a DataTable (in my case have created an widget that inherited the DataTable but i got the same issue with DataTable widget) and it wanted to add rows to the table that had existing rows but also needed to add new columns (which can happen when paginating though a DynamoDB table) it errors out when it tries to render the new table with a column key error. Here is an example app that showcases it:

from textual.app import App, ComposeResult
from textual.widgets import DataTable

class DataTableAppFail(App):
    def compose(self) -> ComposeResult:
        yield DataTable()
    
    def on_mount(self):
        table = self.query_one(DataTable)
        
        table.add_columns(*["col1", "col2", "col3"])
        table.add_rows([[1, 2, 3], [4, 5, 6]])

        table.add_columns(*["col4", "col5"])

DataTableAppFail().run()

when that runs give it gives this error

╭──────────────────────────────────────────────────────── Traceback (most recent call last) ────────────────────────────────────────────────────────╮
│ /home/bob/.cache/pypoetry/virtualenvs/dyno-viewer--bx-3z0J-py3.10/lib/python3.10/site-packages/textual/widgets/_data_table.py:1434 in _on_idle  │
│                                                                                                                                                   │
│   1431 │   │   │   self._require_update_dimensions = False                                                                                        │
│   1432 │   │   │   new_rows = self._new_rows.copy()                                                                                               │
│   1433 │   │   │   self._new_rows.clear()                                                                                                         │
│ ❱ 1434 │   │   │   self._update_dimensions(new_rows)                                                                                              │
│   1435 │   │                                                                                                                                      │
│   1436 │   │   if self._updated_cells:                                                                                                            │
│   1437 │   │   │   # Cell contents have already been updated at this point.                                                                       │
│                                                                                                                                                   │
│ ╭────────────────────── locals ───────────────────────╮                                                                                           │
│ │        _ = Idle()                                   │                                                                                           │
│ │ new_rows = {RowKey(value=None), RowKey(value=None)} │                                                                                           │
│ │     self = DataTable()                              │                                                                                           │
│ ╰─────────────────────────────────────────────────────╯                                                                                           │
│                                                                                                                                                   │
│ /home/bob/.cache/pypoetry/virtualenvs/dyno-viewer--bx-3z0J-py3.10/lib/python3.10/site-packages/textual/widgets/_data_table.py:1144 in           │
│ _update_dimensions                                                                                                                                │
│                                                                                                                                                   │
│   1141 │   │   │   if row.label is not None:                                                                                                      │
│   1142 │   │   │   │   self._labelled_row_exists = True                                                                                           │
│   1143 │   │   │                                                                                                                                  │
│ ❱ 1144 │   │   │   row_label, cells_in_row = self._get_row_renderables(row_index)                                                                 │
│   1145 │   │   │   label_content_width = measure(console, row_label, 1) if row_label else 0                                                       │
│   1146 │   │   │   self._label_column.content_width = max(                                                                                        │
│   1147 │   │   │   │   self._label_column.content_width, label_content_width                                                                      │
│                                                                                                                                                   │
│ ╭─────────────────────────── locals ────────────────────────────╮                                                                                 │
│ │   console = <console width=149 ColorSystem.TRUECOLOR>         │                                                                                 │
│ │  new_rows = {RowKey(value=None), RowKey(value=None)}          │                                                                                 │
│ │       row = Row(key=RowKey(value=None), height=1, label=None) │                                                                                 │
│ │ row_index = 0                                                 │                                                                                 │
│ │   row_key = RowKey(value=None)                                │                                                                                 │
│ │      self = DataTable()                                       │                                                                                 │
│ ╰───────────────────────────────────────────────────────────────╯                                                                                 │
│                                                                                                                                                   │
│ /home/bob/.cache/pypoetry/virtualenvs/dyno-viewer--bx-3z0J-py3.10/lib/python3.10/site-packages/textual/widgets/_data_table.py:1591 in           │
│ _get_row_renderables                                                                                                                              │
│                                                                                                                                                   │
│   1588 │   │   │   # This is the cell where header and row labels intersect                                                                       │
│   1589 │   │   │   return RowRenderables(None, header_row)                                                                                        │
│   1590 │   │                                                                                                                                      │
│ ❱ 1591 │   │   ordered_row = self.get_row_at(row_index)                                                                                           │
│   1592 │   │   empty = Text()                                                                                                                     │
│   1593 │   │                                                                                                                                      │
│   1594 │   │   formatted_row_cells = [                                                                                                            │
│                                                                                                                                                   │
│ ╭────────────────────────────────────────────────────────── locals ───────────────────────────────────────────────────────────╮                   │
│ │ ordered_columns = [                                                                                                         │                   │
│ │                   │   Column(                                                                                               │                   │
│ │                   │   │   key=ColumnKey(value=None),                                                                        │                   │
│ │                   │   │   label=<text 'col1' []>,                                                                           │                   │
│ │                   │   │   width=4,                                                                                          │                   │
│ │                   │   │   content_width=4,                                                                                  │                   │
│ │                   │   │   auto_width=True                                                                                   │                   │
│ │                   │   ),                                                                                                    │                   │
│ │                   │   Column(                                                                                               │                   │
│ │                   │   │   key=ColumnKey(value=None),                                                                        │                   │
│ │                   │   │   label=<text 'col2' []>,                                                                           │                   │
│ │                   │   │   width=4,                                                                                          │                   │
│ │                   │   │   content_width=4,                                                                                  │                   │
│ │                   │   │   auto_width=True                                                                                   │                   │
│ │                   │   ),                                                                                                    │                   │
│ │                   │   Column(                                                                                               │                   │
│ │                   │   │   key=ColumnKey(value=None),                                                                        │                   │
│ │                   │   │   label=<text 'col3' []>,                                                                           │                   │
│ │                   │   │   width=4,                                                                                          │                   │
│ │                   │   │   content_width=4,                                                                                  │                   │
│ │                   │   │   auto_width=True                                                                                   │                   │
│ │                   │   ),                                                                                                    │                   │
│ │                   │   Column(                                                                                               │                   │
│ │                   │   │   key=ColumnKey(value=None),                                                                        │                   │
│ │                   │   │   label=<text 'col4' []>,                                                                           │                   │
│ │                   │   │   width=4,                                                                                          │                   │
│ │                   │   │   content_width=4,                                                                                  │                   │
│ │                   │   │   auto_width=True                                                                                   │                   │
│ │                   │   ),                                                                                                    │                   │
│ │                   │   Column(                                                                                               │                   │
│ │                   │   │   key=ColumnKey(value=None),                                                                        │                   │
│ │                   │   │   label=<text 'col5' []>,                                                                           │                   │
│ │                   │   │   width=4,                                                                                          │                   │
│ │                   │   │   content_width=4,                                                                                  │                   │
│ │                   │   │   auto_width=True                                                                                   │                   │
│ │                   │   )                                                                                                     │                   │
│ │                   ]                                                                                                         │                   │
│ │       row_index = 0                                                                                                         │                   │
│ │            self = DataTable()                                                                                               │                   │
│ ╰─────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────╯                   │
│                                                                                                                                                   │
│ /home/bob/.cache/pypoetry/virtualenvs/dyno-viewer--bx-3z0J-py3.10/lib/python3.10/site-packages/textual/widgets/_data_table.py:845 in get_row_at │
│                                                                                                                                                   │
│    842 │   │   if not self.is_valid_row_index(row_index):                                       ╭──────────── locals ────────────╮                │
│    843 │   │   │   raise RowDoesNotExist(f"Row index {row_index!r} is not valid.")              │ row_index = 0                  │                │
│    844 │   │   row_key = self._row_locations.get_key(row_index)                                 │   row_key = RowKey(value=None) │                │
│ ❱  845 │   │   return self.get_row(row_key)                                                     │      self = DataTable()        │                │
│    846 │                                                                                        ╰────────────────────────────────╯                │
│    847 │   def get_column(self, column_key: ColumnKey | str) -> Iterable[CellType]:                                                               │
│    848 │   │   """Get the values from the column identified by the given column key.                                                              │
│                                                                                                                                                   │
│ /home/bob/.cache/pypoetry/virtualenvs/dyno-viewer--bx-3z0J-py3.10/lib/python3.10/site-packages/textual/widgets/_data_table.py:823 in get_row    │
│                                                                                                                                                   │
│    820 │   │   if row_key not in self._row_locations:                                                                                             │
│    821 │   │   │   raise RowDoesNotExist(f"Row key {row_key!r} is not valid.")                                                                    │
│    822 │   │   cell_mapping: dict[ColumnKey, CellType] = self._data.get(row_key, {})                                                              │
│ ❱  823 │   │   ordered_row: list[CellType] = [                                                                                                    │
│    824 │   │   │   cell_mapping[column.key] for column in self.ordered_columns                                                                    │
│    825 │   │   ]                                                                                                                                  │
│    826 │   │   return ordered_row                                                                                                                 │
│                                                                                                                                                   │
│ ╭──────────────────────────────────────── locals ─────────────────────────────────────────╮                                                       │
│ │ cell_mapping = {                                                                        │                                                       │
│ │                │   <textual.widgets._data_table.ColumnKey object at 0x7fa1a8c1bc70>: 1, │                                                       │
│ │                │   <textual.widgets._data_table.ColumnKey object at 0x7fa1a8242890>: 2, │                                                       │
│ │                │   <textual.widgets._data_table.ColumnKey object at 0x7fa1a8242a10>: 3  │                                                       │
│ │                }                                                                        │                                                       │
│ │      row_key = RowKey(value=None)                                                       │                                                       │
│ │         self = DataTable()                                                              │                                                       │
│ ╰─────────────────────────────────────────────────────────────────────────────────────────╯                                                       │
│                                                                                                                                                   │
│ /home/bob/.cache/pypoetry/virtualenvs/dyno-viewer--bx-3z0J-py3.10/lib/python3.10/site-packages/textual/widgets/_data_table.py:824 in <listcomp> │
│                                                                                                                                                   │
│    821 │   │   │   raise RowDoesNotExist(f"Row key {row_key!r} is not valid.")                                                                    │
│    822 │   │   cell_mapping: dict[ColumnKey, CellType] = self._data.get(row_key, {})                                                              │
│    823 │   │   ordered_row: list[CellType] = [                                                                                                    │
│ ❱  824 │   │   │   cell_mapping[column.key] for column in self.ordered_columns                                                                    │
│    825 │   │   ]                                                                                                                                  │
│    826 │   │   return ordered_row                                                                                                                 │
│    827                                                                                                                                            │
│                                                                                                                                                   │
│ ╭────────────────────────────────────────────────────── locals ───────────────────────────────────────────────────────╮                           │
│ │           .0 = <list_iterator object at 0x7fa1a8279a80>                                                             │                           │
│ │ cell_mapping = {                                                                                                    │                           │
│ │                │   <textual.widgets._data_table.ColumnKey object at 0x7fa1a8c1bc70>: 1,                             │                           │
│ │                │   <textual.widgets._data_table.ColumnKey object at 0x7fa1a8242890>: 2,                             │                           │
│ │                │   <textual.widgets._data_table.ColumnKey object at 0x7fa1a8242a10>: 3                              │                           │
│ │                }                                                                                                    │                           │
│ │       column = Column(key=ColumnKey(value=None), label=<text 'col4' []>, width=4, content_width=4, auto_width=True) │                           │
│ ╰─────────────────────────────────────────────────────────────────────────────────────────────────────────────────────╯                           │
╰───────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────╯
KeyError: <textual.widgets._data_table.ColumnKey object at 0x7fa1a8240730>

This looks like a bug but wanted to check if i was doing something wrong or if you have any suggestions

Thanks

[TomJGooding]

I agree this looks like a bug. Perhaps add_column should update the table data if necessary to add None values?

If you clear the table and add back the same rows, it works as expected thanks to itertools.zip_longest.

[mrllama123]

yeah that's what i am thinking as well. I did something similar to that to work around this issue but i feel like this should be part of the DataTable class

[TomJGooding]

@mrllama123 Fixed by @darrenburns in the latest version!