From 98b9bd1b12362726ae803729afa5fd079705cea1 Mon Sep 17 00:00:00 2001 From: harshitgupta31415 Date: Sun, 20 Sep 2026 11:23:23 +0530 Subject: [PATCH 1/5] Optimize Dataset.subset() lookups --- AUTHORS | 1 + src/tablib/core.py | 12 +++--------- tests/test_tablib.py | 15 +++++++++++++++ 3 files changed, 19 insertions(+), 9 deletions(-) diff --git a/AUTHORS b/AUTHORS index cd73e657..f143a515 100644 --- a/AUTHORS +++ b/AUTHORS @@ -12,6 +12,7 @@ Here is a list of past and present much-appreciated contributors: Daniel Santos Egor Osokin Erik Youngren + Harshit Gupta Hugo van Kemenade Ian Stride Iuri de Silvio diff --git a/src/tablib/core.py b/src/tablib/core.py index b0b56fd2..6dfab396 100644 --- a/src/tablib/core.py +++ b/src/tablib/core.py @@ -835,8 +835,9 @@ def subset(self, rows=None, cols=None): cols = list(self.headers) # filter out impossible rows and columns - rows = [row for row in rows if row in range(self.height)] + rows = {row for row in rows if row in range(self.height)} cols = [header for header in cols if header in self.headers] + col_indexes = [self.headers.index(header) for header in cols] _dset = Dataset() @@ -845,15 +846,8 @@ def subset(self, rows=None, cols=None): _dset._data = [] for row_no, row in enumerate(self._data): - data_row = [] - for key in _dset.headers: - if key in self.headers: - pos = self.headers.index(key) - data_row.append(row[pos]) - else: - raise KeyError - if row_no in rows: + data_row = [row[index] for index in col_indexes] _dset.append(row=Row(data_row)) return _dset diff --git a/tests/test_tablib.py b/tests/test_tablib.py index ca8ac05e..ad3a5f80 100755 --- a/tests/test_tablib.py +++ b/tests/test_tablib.py @@ -610,6 +610,21 @@ def test_subset(self): self.assertEqual(subset._data[0].list, ['John', 90]) self.assertEqual(subset._data[1].list, ['Thomas', 50]) + def test_subset_preserves_selection_behavior(self): + data.headers = self.headers + data.append(self.john) + data.append(self.george) + data.append(self.tom) + + subset = data.subset( + rows=(2, 0, 2, -1, 99), + cols=('gpa', 'missing', 'first_name', 'gpa'), + ) + + self.assertEqual(subset.headers, ['gpa', 'first_name', 'gpa']) + self.assertEqual(subset._data[0].list, [90, 'John', 90]) + self.assertEqual(subset._data[1].list, [50, 'Thomas', 50]) + def test_formatters(self): """Confirm formatters are being triggered.""" From bf2a3073f45209d5bb0d17c21ef7f2c0207f8e4f Mon Sep 17 00:00:00 2001 From: harshitgupta31415 Date: Sun, 20 Sep 2026 22:44:01 +0530 Subject: [PATCH 2/5] Preserve subset behavior without valid columns --- src/tablib/core.py | 2 +- tests/test_tablib.py | 9 +++++++++ 2 files changed, 10 insertions(+), 1 deletion(-) diff --git a/src/tablib/core.py b/src/tablib/core.py index 6dfab396..2e3a75e7 100644 --- a/src/tablib/core.py +++ b/src/tablib/core.py @@ -837,12 +837,12 @@ def subset(self, rows=None, cols=None): # filter out impossible rows and columns rows = {row for row in rows if row in range(self.height)} cols = [header for header in cols if header in self.headers] - col_indexes = [self.headers.index(header) for header in cols] _dset = Dataset() # filtering rows and columns _dset.headers = list(cols) + col_indexes = [self.headers.index(header) for header in _dset.headers] _dset._data = [] for row_no, row in enumerate(self._data): diff --git a/tests/test_tablib.py b/tests/test_tablib.py index ad3a5f80..cdfc1fc0 100755 --- a/tests/test_tablib.py +++ b/tests/test_tablib.py @@ -625,6 +625,15 @@ def test_subset_preserves_selection_behavior(self): self.assertEqual(subset._data[0].list, [90, 'John', 90]) self.assertEqual(subset._data[1].list, [50, 'Thomas', 50]) + def test_subset_preserves_no_matching_column_behavior(self): + data.headers = self.headers + data.append(self.john) + + for columns in ([], ['missing']): + with self.subTest(columns=columns): + with self.assertRaisesRegex(TypeError, "'NoneType' object is not iterable"): + data.subset(cols=columns) + def test_formatters(self): """Confirm formatters are being triggered.""" From 4357722e8eb30260458d6e9b29fafc5e1bd9ace7 Mon Sep 17 00:00:00 2001 From: harshitgupta31415 Date: Fri, 25 Sep 2026 08:57:55 +0530 Subject: [PATCH 3/5] Preserve subset selection for custom row indexes --- src/tablib/core.py | 5 ++++- tests/test_tablib.py | 16 ++++++++++++++++ 2 files changed, 20 insertions(+), 1 deletion(-) diff --git a/src/tablib/core.py b/src/tablib/core.py index 2e3a75e7..9a50a13c 100644 --- a/src/tablib/core.py +++ b/src/tablib/core.py @@ -835,7 +835,10 @@ def subset(self, rows=None, cols=None): cols = list(self.headers) # filter out impossible rows and columns - rows = {row for row in rows if row in range(self.height)} + rows = [row for row in rows if row in range(self.height)] + # Keep list membership for selectors with custom equality or hashing. + if all(type(row) is int for row in rows): + rows = set(rows) cols = [header for header in cols if header in self.headers] _dset = Dataset() diff --git a/tests/test_tablib.py b/tests/test_tablib.py index cdfc1fc0..117d1979 100755 --- a/tests/test_tablib.py +++ b/tests/test_tablib.py @@ -634,6 +634,22 @@ def test_subset_preserves_no_matching_column_behavior(self): with self.assertRaisesRegex(TypeError, "'NoneType' object is not iterable"): data.subset(cols=columns) + def test_subset_preserves_unhashable_row_selection(self): + class FirstRow: + __hash__ = None + + def __eq__(self, other): + return other == 0 + + data.headers = self.headers + data.append(self.john) + data.append(self.george) + + subset = data.subset(rows=[1, FirstRow()], cols=['first_name']) + self.assertEqual(subset[0], ('John',)) + self.assertEqual(subset[1], ('George',)) + self.assertEqual(subset.height, 2) + def test_formatters(self): """Confirm formatters are being triggered.""" From 8e982316b915750dfd2fb0b0dea9aba8bf2cc5ab Mon Sep 17 00:00:00 2001 From: harshitgupta31415 Date: Sun, 27 Sep 2026 14:07:17 +0530 Subject: [PATCH 4/5] Preserve subset column errors and custom header comparisons --- src/tablib/core.py | 27 ++++++- tests/test_subset_compatibility.py | 116 +++++++++++++++++++++++++++++ 2 files changed, 141 insertions(+), 2 deletions(-) create mode 100644 tests/test_subset_compatibility.py diff --git a/src/tablib/core.py b/src/tablib/core.py index 9a50a13c..2b994707 100644 --- a/src/tablib/core.py +++ b/src/tablib/core.py @@ -845,9 +845,32 @@ def subset(self, rows=None, cols=None): # filtering rows and columns _dset.headers = list(cols) - col_indexes = [self.headers.index(header) for header in _dset.headers] - _dset._data = [] + # Custom headers can change equality or raise their own exceptions. + # Keep their original per-row comparisons instead of caching them. + if ( + type(self) is not Dataset + or not all(type(header) is str for header in self.headers) + or not all(type(header) is str for header in _dset.headers) + ): + for row_no, row in enumerate(self._data): + data_row = [] + for header in _dset.headers: + if header in self.headers: + data_row.append(row[self.headers.index(header)]) + else: + raise KeyError + if row_no in rows: + _dset.append(row=Row(data_row)) + return _dset + + col_indexes = [] + for header in _dset.headers: + # A selector may have changed headers while it was consumed. + if header not in self.headers: + raise KeyError + col_indexes.append(self.headers.index(header)) + for row_no, row in enumerate(self._data): if row_no in rows: data_row = [row[index] for index in col_indexes] diff --git a/tests/test_subset_compatibility.py b/tests/test_subset_compatibility.py new file mode 100644 index 00000000..7dee8d58 --- /dev/null +++ b/tests/test_subset_compatibility.py @@ -0,0 +1,116 @@ +from copy import copy + +import pytest + +from tablib import Dataset +from tablib.core import Row + + +def original_subset(dataset, rows=None, cols=None): + if not dataset: + return + if rows is None: + rows = list(range(dataset.height)) + if cols is None: + cols = list(dataset.headers) + rows = [row for row in rows if row in range(dataset.height)] + cols = [header for header in cols if header in dataset.headers] + result = Dataset() + result.headers = list(cols) + result._data = [] + for row_no, row in enumerate(dataset._data): + data_row = [] + for header in result.headers: + if header in dataset.headers: + data_row.append(row[dataset.headers.index(header)]) + else: + raise KeyError + if row_no in rows: + result.append(row=Row(data_row)) + return result + + +@pytest.mark.parametrize('rows', [None, [], [0]]) +def test_subset_selector_changing_headers_keeps_key_error(rows): + dataset = Dataset(['Synthetic A', 20], ['Synthetic B', 30], headers=['name', 'age']) + + def columns(): + yield 'name' + dataset.headers = ['renamed', 'age'] + + with pytest.raises(KeyError) as error: + dataset.subset(rows=rows, cols=columns()) + assert error.value.args == () + assert error.value.__cause__ is None + assert error.value.__context__ is None + + +@pytest.mark.parametrize('rows', [None, [], [0]]) +@pytest.mark.parametrize('at,raises', [(2, False), (4, False), (2, True), (3, True), (4, True)]) +def test_subset_preserves_stateful_comparisons_and_their_errors(rows, at, raises): + def run(method): + calls = [] + + class Column(str): + def __eq__(self, other): + calls.append(other) + if len(calls) == at: + if raises: + raise ValueError('custom comparison error') + return False + return super().__eq__(other) + + dataset = Dataset(['Synthetic A'], ['Synthetic B'], headers=['name']) + try: + result = method(dataset, rows=rows, cols=[Column('name')]) + outcome = list(result) + except (KeyError, ValueError) as error: + outcome = type(error), error.args + return outcome, calls + + assert run(Dataset.subset) == run(original_subset) + + +@pytest.mark.parametrize('headers,columns', [ + ([1, 2], [2, 1]), + ([None, 'age'], [None]), + (['name', 'name'], ['name', 'name']), + (['name', 'age'], ['missing', 'name', 'name']), +]) +def test_subset_nonstring_and_duplicate_headers_match_original(headers, columns): + dataset = Dataset(['Synthetic A', 20], ['Synthetic B', 30], headers=headers) + expected = original_subset(dataset, cols=columns) + actual = dataset.subset(cols=columns) + assert actual.headers == expected.headers + assert list(actual) == list(expected) + + +def test_subset_subclass_keeps_header_access_behavior(): + class ObservedDataset(Dataset): + def __init__(self, *args, **kwargs): + self.header_reads = 0 + super().__init__(*args, **kwargs) + + @property + def headers(self): + self.header_reads += 1 + return Dataset.headers.fget(self) + + @headers.setter + def headers(self, collection): + Dataset.headers.fset(self, collection) + + def run(method): + dataset = ObservedDataset(['Synthetic A'], ['Synthetic B'], headers=['name']) + dataset.header_reads = 0 + result = method(dataset, rows=[0], cols=['name']) + return list(result), dataset.header_reads + + assert run(Dataset.subset) == run(original_subset) + + +def test_subset_does_not_change_the_input(): + dataset = Dataset(['Synthetic A', 20], headers=['name', 'age']) + state = copy(dataset._data), dataset.headers[:] + dataset.subset(cols=['age', 'name', 'age']) + assert (dataset._data, dataset.headers) == state From 78ecabd3e3089f0a67623498e974d7591c8e4f11 Mon Sep 17 00:00:00 2001 From: harshitgupta31415 Date: Sun, 27 Sep 2026 15:19:57 +0530 Subject: [PATCH 5/5] Preserve subset edge cases and avoid whole-header scans --- src/tablib/core.py | 55 +++++++-- tests/test_subset_compatibility.py | 172 ++++++++++++++++++++++++++++- 2 files changed, 211 insertions(+), 16 deletions(-) diff --git a/src/tablib/core.py b/src/tablib/core.py index 2b994707..a62d6450 100644 --- a/src/tablib/core.py +++ b/src/tablib/core.py @@ -107,6 +107,9 @@ def has_tag(self, tag): return bool(len(set(tag) & set(self.tags))) +_DEFAULT_ROW_GETITEM = Row.__getitem__ + + class Dataset: """The :class:`Dataset` object is the heart of Tablib. It provides all core functionality. @@ -846,13 +849,48 @@ def subset(self, rows=None, cols=None): # filtering rows and columns _dset.headers = list(cols) _dset._data = [] - # Custom headers can change equality or raise their own exceptions. - # Keep their original per-row comparisons instead of caching them. + col_indexes = None + header_width = ( + len(self.headers) + if type(self) is Dataset and type(self.headers) is list + else None + ) + # Small tables and custom/inconsistent inputs keep the original path, + # including reads of unselected rows and their possible exceptions. if ( - type(self) is not Dataset - or not all(type(header) is str for header in self.headers) - or not all(type(header) is str for header in _dset.headers) + type(self) is Dataset + and type(self._data) is list + and len(self._data) >= 16 + and type(rows) is set + and Row.__getitem__ is _DEFAULT_ROW_GETITEM + and header_width is not None + and all(type(header) is str for header in _dset.headers) + and all( + type(row) is Row + and type(row._row) is list + and len(row._row) == header_width + for row in self._data + ) ): + # Stop at the last needed first occurrence, not the end of a wide + # header list. Never compare a custom header while probing: let + # the original path perform its comparisons and side effects. + remaining = set(_dset.headers) + positions = {} + for index, header in enumerate(self.headers): + if type(header) is not str: + break + if header in remaining: + positions[header] = index + remaining.remove(header) + if not remaining: + col_indexes = [positions[header] for header in _dset.headers] + break + else: + # A selector may have removed a previously accepted header. + raise KeyError + + if col_indexes is None: for row_no, row in enumerate(self._data): data_row = [] for header in _dset.headers: @@ -864,13 +902,6 @@ def subset(self, rows=None, cols=None): _dset.append(row=Row(data_row)) return _dset - col_indexes = [] - for header in _dset.headers: - # A selector may have changed headers while it was consumed. - if header not in self.headers: - raise KeyError - col_indexes.append(self.headers.index(header)) - for row_no, row in enumerate(self._data): if row_no in rows: data_row = [row[index] for index in col_indexes] diff --git a/tests/test_subset_compatibility.py b/tests/test_subset_compatibility.py index 7dee8d58..c649be36 100644 --- a/tests/test_subset_compatibility.py +++ b/tests/test_subset_compatibility.py @@ -1,4 +1,5 @@ from copy import copy +from unittest.mock import patch import pytest @@ -31,8 +32,9 @@ def original_subset(dataset, rows=None, cols=None): @pytest.mark.parametrize('rows', [None, [], [0]]) -def test_subset_selector_changing_headers_keeps_key_error(rows): - dataset = Dataset(['Synthetic A', 20], ['Synthetic B', 30], headers=['name', 'age']) +@pytest.mark.parametrize('height', [2, 16]) +def test_subset_selector_changing_headers_keeps_key_error(rows, height): + dataset = Dataset(*(['Synthetic A', 20] for _ in range(height)), headers=['name', 'age']) def columns(): yield 'name' @@ -47,7 +49,8 @@ def columns(): @pytest.mark.parametrize('rows', [None, [], [0]]) @pytest.mark.parametrize('at,raises', [(2, False), (4, False), (2, True), (3, True), (4, True)]) -def test_subset_preserves_stateful_comparisons_and_their_errors(rows, at, raises): +@pytest.mark.parametrize('height', [2, 16]) +def test_subset_preserves_stateful_comparisons_and_their_errors(rows, at, raises, height): def run(method): calls = [] @@ -60,7 +63,7 @@ def __eq__(self, other): return False return super().__eq__(other) - dataset = Dataset(['Synthetic A'], ['Synthetic B'], headers=['name']) + dataset = Dataset(*(['Synthetic A'] for _ in range(height)), headers=['name']) try: result = method(dataset, rows=rows, cols=[Column('name')]) outcome = list(result) @@ -114,3 +117,164 @@ def test_subset_does_not_change_the_input(): state = copy(dataset._data), dataset.headers[:] dataset.subset(cols=['age', 'name', 'age']) assert (dataset._data, dataset.headers) == state + + +@pytest.mark.parametrize('height', [2, 16, 64]) +@pytest.mark.parametrize('action', ['wipe', 'remove_and_rename', 'rows_with_no_columns']) +def test_subset_selector_emptying_the_source_matches_original(height, action): + def run(method): + dataset = Dataset(*([index] for index in range(height)), headers=['value']) + + def columns(): + yield 'value' + if action == 'wipe': + dataset.wipe() + else: + del dataset[:] + dataset.headers = ['renamed'] + + def rows(): + del dataset[:] + yield 0 + + if action == 'rows_with_no_columns': + result = method(dataset, rows=rows(), cols=[]) + else: + result = method(dataset, cols=columns()) + return list(result), result.headers + + assert run(Dataset.subset) == run(original_subset) + + +@pytest.mark.parametrize('height', [2, 16, 64]) +@pytest.mark.parametrize('rows', [[], [0], [1]]) +def test_subset_keeps_errors_from_unselected_short_rows(height, rows): + def run(method): + dataset = Dataset(*([index, index + 100] for index in range(height - 1)), + [999], headers=['a', 'b']) + with pytest.raises(IndexError) as error: + method(dataset, rows=rows, cols=['b']) + return error.value.args + + assert run(Dataset.subset) == run(original_subset) + + +@pytest.mark.parametrize('height', [2, 16, 64]) +@pytest.mark.parametrize('rows', [[], [0]]) +def test_subset_keeps_index_errors_after_public_header_mutation(height, rows): + dataset = Dataset(*([index, index + 100] for index in range(height)), headers=['a', 'b']) + dataset.headers.append('extra') + with pytest.raises(IndexError) as expected: + original_subset(dataset, rows=rows, cols=['extra']) + with pytest.raises(IndexError) as actual: + dataset.subset(rows=rows, cols=['extra']) + assert actual.value.args == expected.value.args + + +@pytest.mark.parametrize('height', [2, 16, 64]) +@pytest.mark.parametrize('customize', ['row', 'cells', 'container', 'method']) +def test_subset_keeps_custom_row_reads_and_header_changes(height, customize): + def run(method): + dataset = Dataset(*([index, index + 100] for index in range(height)), + headers=['a', 'b']) + events = [] + original_getitem = Row.__getitem__ + + def read(row, index): + value = original_getitem(row, index) + events.append(index) + dataset.headers[:] = ['b', 'a'] + return value + + class CustomRow(Row): + def __getitem__(self, index): + return read(self, index) + + class Cells(list): + def __getitem__(self, index): + value = super().__getitem__(index) + events.append(index) + dataset.headers[:] = ['b', 'a'] + return value + + class Rows(list): + def __iter__(self): + events.append('iter') + dataset.headers[:] = ['b', 'a'] + return super().__iter__() + + if customize == 'row': + dataset._data[0] = CustomRow([0, 100]) + elif customize == 'cells': + dataset._data[0]._row = Cells([0, 100]) + elif customize == 'container': + dataset._data = Rows(dataset._data) + if customize == 'method': + with patch.object(Row, '__getitem__', read): + result = method(dataset, cols=['a']) + else: + result = method(dataset, cols=['a']) + return list(result), dataset.headers, events + + assert run(Dataset.subset) == run(original_subset) + + +def test_subset_keeps_selector_comparison_side_effects(): + def run(method): + dataset = Dataset(*([index, index + 100] for index in range(16)), + headers=['a', 'b']) + calls = [] + + class Selector: + __hash__ = None + + def __eq__(self, other): + calls.append(other) + if len(calls) > 1: + dataset.headers.reverse() + return True + + result = method(dataset, rows=[Selector()], cols=['a']) + return list(result), dataset.headers, calls + + assert run(Dataset.subset) == run(original_subset) + + +@pytest.mark.parametrize('custom_first', [True, False]) +def test_subset_source_header_comparison_calls_match_original(custom_first): + def run(method): + calls = [] + + class Header(str): + def __eq__(self, other): + calls.append(other) + return super().__eq__(other) + + headers = [Header('custom'), 'wanted'] if custom_first else ['wanted', Header('custom')] + dataset = Dataset(*([index, index + 100] for index in range(16)), headers=headers) + result = method(dataset, rows=[0], cols=['wanted']) + return list(result), calls + + assert run(Dataset.subset) == run(original_subset) + + +def test_subset_uses_first_header_occurrences_on_the_cached_path(): + dataset = Dataset(*([index, index + 100, index + 200] for index in range(16)), + headers=['a', 'a', 'b']) + expected = original_subset(dataset, rows=[15, 0, 15], cols=['b', 'a', 'b']) + result = dataset.subset(rows=[15, 0, 15], cols=['b', 'a', 'b']) + assert result.headers == expected.headers + assert list(result) == list(expected) + + +def test_subset_header_removed_by_selector_keeps_key_error_on_cached_path(): + dataset = Dataset(*([index, index + 100] for index in range(16)), headers=['a', 'b']) + + def columns(): + yield 'a' + dataset.headers = ['renamed', 'b'] + + with pytest.raises(KeyError) as error: + dataset.subset(cols=columns()) + assert error.value.args == () + assert error.value.__context__ is None