From 01ac5e6dc60407f6dae7c4b53ff95ae840a49a17 Mon Sep 17 00:00:00 2001
From: Claude Paroz <claude@2xlibre.net>
Date: Thu, 21 Dec 2023 08:27:18 +0100
Subject: [PATCH] Fixes #572 - Allow appending rows after a dynamic column was
 inserted

Co-authored-by: Hugo van Kemenade <hugovk@users.noreply.github.com>
---
 HISTORY.md           |  3 +++
 docs/tutorial.rst    |  9 +++++++++
 src/tablib/core.py   | 22 +++++++++++++++++++---
 tests/test_tablib.py | 36 ++++++++++++++++++++++++++++++++++--
 4 files changed, 65 insertions(+), 5 deletions(-)

diff --git a/HISTORY.md b/HISTORY.md
index 17551004..71350af4 100644
--- a/HISTORY.md
+++ b/HISTORY.md
@@ -7,6 +7,9 @@
 - The html format now supports importing from HTML content (#243)
 - The ODS format now supports importing from .ods files (#567). The support is
   still a bit experimental.
+- When adding rows to a dataset with dynamic columns, it's now possible to
+  provide only static values, and dynamic column values will be automatically
+  calculated and added to the row (#572).
 
 ### Changes
 
diff --git a/docs/tutorial.rst b/docs/tutorial.rst
index e5cd12bc..d21e7fcb 100644
--- a/docs/tutorial.rst
+++ b/docs/tutorial.rst
@@ -278,6 +278,15 @@ Adding this function to our dataset as a dynamic column would result in: ::
     - {Age: 22, First Name: Kenneth, Gender: Male, Last Name: Reitz}
     - {Age: 20, First Name: Bessie, Gender: Female, Last Name: Monke}
 
+When you add new rows to a dataset that contains dynamic columns, you should
+either provide all values in the row, or only the non-dynamic values and then
+the dynamic values will be automatically generated using the function initially
+provided for the column calculation.
+
+..versionchanged:: 3.6.0
+
+    In older versions, you could only add new rows with fully-populated rows,
+    including dynamic columns.
 
 .. _tags:
 
diff --git a/src/tablib/core.py b/src/tablib/core.py
index a7ebff2b..ef3f6fe2 100644
--- a/src/tablib/core.py
+++ b/src/tablib/core.py
@@ -155,6 +155,9 @@ def __init__(self, *args, **kwargs):
         # (column, callback) tuples
         self._formatters = []
 
+        # {col_index: col_func}
+        self._dynamic_columns = {}
+
         self.headers = kwargs.get('headers')
 
         self.title = kwargs.get('title')
@@ -187,6 +190,8 @@ def __delitem__(self, key):
 
                 pos = self.headers.index(key)
                 del self.headers[pos]
+                if pos in self._dynamic_columns:
+                    del self._dynamic_columns[pos]
 
                 for i, row in enumerate(self._data):
 
@@ -238,7 +243,13 @@ def _set_in_format(self, fmt_key, in_stream, **kwargs):
     def _validate(self, row=None, col=None, safety=False):
         """Assures size of every row in dataset is of proper proportions."""
         if row:
-            is_valid = (len(row) == self.width) if self.width else True
+            if self.width:
+                is_valid = (
+                    len(row) == self.width or
+                    len(row) == (self.width - len(self._dynamic_columns))
+                )
+            else:
+                is_valid = True
         elif col:
             if len(col) < 1:
                 is_valid = True
@@ -446,9 +457,13 @@ def insert(self, index, row, tags=()):
 
         The default behaviour is to insert the given row to the :class:`Dataset`
         object at the given index.
-       """
+        """
 
         self._validate(row)
+        if len(row) < self.width:
+            for pos, func in self._dynamic_columns.items():
+                row = list(row)
+                row.insert(pos, func(row))
         self._data.insert(index, Row(row, tags=tags))
 
     def rpush(self, row, tags=()):
@@ -546,7 +561,8 @@ def insert_col(self, index, col=None, header=None):
             col = []
 
         # Callable Columns...
-        if hasattr(col, '__call__'):
+        if callable(col):
+            self._dynamic_columns[self.width] = col
             col = list(map(col, self._data))
 
         col = self._clean_col(col)
diff --git a/tests/test_tablib.py b/tests/test_tablib.py
index 9add226e..9035c1e1 100755
--- a/tests/test_tablib.py
+++ b/tests/test_tablib.py
@@ -181,10 +181,42 @@ def test_add_column_with_header_and_data_exists(self):
     def test_add_callable_column(self):
         """Verify adding column with values specified as callable."""
 
-        def new_col(x):
-            return x[0]
+        def new_col(row):
+            return row[0]
+
+        def initials(row):
+            return f"{row[0][0]}{row[1][0]}"
 
         self.founders.append_col(new_col, header='first_again')
+        self.founders.append_col(initials, header='initials')
+
+        # A new row can still be appended, and the dynamic column value generated.
+        self.founders.append(('Some', 'One', 71))
+        # Also acceptable when all dynamic column values are provided.
+        self.founders.append(('Other', 'Second', 84, 'Other', 'OS'))
+
+        self.assertEqual(self.founders[3], ('Some', 'One', 71, 'Some', 'SO'))
+        self.assertEqual(self.founders[4], ('Other', 'Second', 84, 'Other', 'OS'))
+        self.assertEqual(
+            self.founders['first_again'],
+            ['John', 'George', 'Thomas', 'Some', 'Other']
+        )
+        self.assertEqual(
+            self.founders['initials'],
+            ['JA', 'GW', 'TJ', 'SO', 'OS']
+        )
+
+        # However only partial dynamic values provided is not accepted.
+        with self.assertRaises(tablib.InvalidDimensions):
+            self.founders.append(('Should', 'Crash', 60, 'Partial'))
+
+        # Add a new row after dynamic column deletion
+        del self.founders['first_again']
+        self.founders.append(('After', 'Deletion', 75))
+        self.assertEqual(
+            self.founders['initials'],
+            ['JA', 'GW', 'TJ', 'SO', 'OS', 'AD']
+        )
 
     def test_header_slicing(self):
         """Verify slicing by headers."""