Fix DataTable crash when removing and updating cell at the same time

CHANGELOG.md

@@ -15,6 +15,7 @@ and this project adheres to [Semantic Versioning](http://semver.org/).
 - Fix location of IME and emoji popups https://github.com/Textualize/textual/pull/3408
 - Fixed application freeze when pasting an emoji into an application on Windows https://github.com/Textualize/textual/issues/3178
 - Fixed duplicate option ID handling in the `OptionList` https://github.com/Textualize/textual/issues/3455
+- Fix crash when removing and updating DataTable cell at same time https://github.com/Textualize/textual/pull/3487
 - Fixed fractional styles to allow integer values https://github.com/Textualize/textual/issues/3414
 
 ### Added

src/textual/widgets/_data_table.py

@@ -1632,6 +1632,10 @@ def remove_row(self, row_key: RowKey | str) -> None:
 
         self._row_locations = new_row_locations
 
+        # Prevent the removed cells from triggering dimension updates
+        for column_key in self._data.get(row_key):
+            self._updated_cells.discard(CellKey(row_key, column_key))
+
         del self.rows[row_key]
         del self._data[row_key]
 
@@ -1668,8 +1672,10 @@ def remove_column(self, column_key: ColumnKey | str) -> None:
         self._column_locations = new_column_locations
 
         del self.columns[column_key]
-        for row in self._data:
-            del self._data[row][column_key]
+
+        for row_key in self._data:
+            self._updated_cells.discard(CellKey(row_key, column_key))
+            del self._data[row_key][column_key]
 
         self.cursor_coordinate = self.cursor_coordinate
         self.hover_coordinate = self.hover_coordinate

tests/test_data_table.py

@@ -309,6 +309,19 @@ async def test_remove_row():
         assert len(table.rows) == 2
 
 
+async def test_remove_row_and_update():
+    """Regression test for https://github.com/Textualize/textual/issues/3470 -
+    Crash when attempting to remove and update the same cell."""
+    app = DataTableApp()
+    async with app.run_test() as pilot:
+        table: DataTable = app.query_one(DataTable)
+        table.add_column("A", key="A")
+        table.add_row("1", key="1")
+        table.update_cell("1", "A", "X", update_width=True)
+        table.remove_row("1")
+        await pilot.pause()
+
+
 async def test_remove_column():
     app = DataTableApp()
     async with app.run_test():
@@ -323,6 +336,19 @@ async def test_remove_column():
         assert table.get_row_at(2) == ["2/1"]
 
 
+async def test_remove_column_and_update():
+    """Regression test for https://github.com/Textualize/textual/issues/3470 -
+    Crash when attempting to remove and update the same cell."""
+    app = DataTableApp()
+    async with app.run_test() as pilot:
+        table: DataTable = app.query_one(DataTable)
+        table.add_column("A", key="A")
+        table.add_row("1", key="1")
+        table.update_cell("1", "A", "X", update_width=True)
+        table.remove_column("A")
+        await pilot.pause()
+
+
 async def test_clear():
     app = DataTableApp()
     async with app.run_test():