diff --git a/src/tablib/core.py b/src/tablib/core.py index b0b56fd2..666f9728 100644 --- a/src/tablib/core.py +++ b/src/tablib/core.py @@ -638,7 +638,7 @@ def add_formatter(self, col, handler): else: raise KeyError - if col is None or col <= self.width: + if col is None or col < self.width: self._formatters.append((col, handler)) else: raise InvalidDatasetIndex diff --git a/tests/test_tablib.py b/tests/test_tablib.py index ca8ac05e..44aaf28f 100755 --- a/tests/test_tablib.py +++ b/tests/test_tablib.py @@ -19,7 +19,7 @@ import tablib from tablib.core import Row, detect_format -from tablib.exceptions import UnsupportedFormat +from tablib.exceptions import InvalidDatasetIndex, UnsupportedFormat from tablib.formats import registry try: @@ -643,6 +643,21 @@ def _formatter(cell_value): {'first_name': 'THOMAS', 'last_name': 'JEFFERSON', 'gpa': '50'}, ]) + def test_add_formatter_out_of_range_column(self): + """A column index past the last column must be rejected, not silently + accepted (which then crashed on export).""" + + def _formatter(cell_value): + return cell_value + + # width is 3 -> valid indices are 0, 1, 2 + self.assertEqual(self.founders.width, 3) + # the last valid index is accepted + self.assertTrue(self.founders.add_formatter(2, _formatter)) + # one past the last column raises instead of being silently accepted + with self.assertRaises(InvalidDatasetIndex): + self.founders.add_formatter(3, _formatter) + def test_unicode_renders_markdown_table(self): # add another entry to test right field width for # integer