diff --git a/tabulator/stream.py b/tabulator/stream.py index 9d96cb86..fcc78851 100644 --- a/tabulator/stream.py +++ b/tabulator/stream.py @@ -752,11 +752,13 @@ def __extract_headers(self): ignore = True # Ignore listed headers if self.__ignore_listed_headers is not None: - if header in self.__ignore_listed_headers: + if (header in self.__ignore_listed_headers or + index + 1 in self.__ignore_listed_headers): ignore = True # Ignore not-listed headers if self.__ignore_not_listed_headers is not None: - if header not in self.__ignore_not_listed_headers: + if (header not in self.__ignore_not_listed_headers and + index + 1 not in self.__ignore_not_listed_headers): ignore = True # Add to the list and skip if ignore: diff --git a/tests/test_stream.py b/tests/test_stream.py index 6abaa92f..63afd9b3 100644 --- a/tests/test_stream.py +++ b/tests/test_stream.py @@ -358,6 +358,26 @@ def test_stream_skip_fields(): ] +def test_stream_skip_fields_position(): + source = 'text://header1,header2,header3\nvalue1,value2,value3' + with Stream(source, format='csv', headers=1, skip_fields=[2]) as stream: + assert stream.headers == ['header1', 'header3'] + assert stream.field_positions == [1, 3] + assert stream.read(keyed=True) == [ + {'header1': 'value1', 'header3': 'value3'}, + ] + + +def test_stream_skip_fields_position_and_prefix(): + source = 'text://header1,header2,header3\nvalue1,value2,value3' + with Stream(source, format='csv', headers=1, skip_fields=[2, 'header3']) as stream: + assert stream.headers == ['header1'] + assert stream.field_positions == [1] + assert stream.read(keyed=True) == [ + {'header1': 'value1'}, + ] + + def test_stream_skip_fields_blank_header(): source = 'text://header1,,header3\nvalue1,value2,value3' with Stream(source, format='csv', headers=1, skip_fields=['']) as stream: @@ -378,6 +398,26 @@ def test_stream_pick_fields(): ] +def test_stream_pick_fields_position(): + source = 'text://header1,header2,header3\nvalue1,value2,value3' + with Stream(source, format='csv', headers=1, pick_fields=[2]) as stream: + assert stream.headers == ['header2'] + assert stream.field_positions == [2] + assert stream.read(keyed=True) == [ + {'header2': 'value2'}, + ] + + +def test_stream_pick_fields_position_and_prefix(): + source = 'text://header1,header2,header3\nvalue1,value2,value3' + with Stream(source, format='csv', headers=1, pick_fields=[2, 'header3']) as stream: + assert stream.headers == ['header2', 'header3'] + assert stream.field_positions == [2, 3] + assert stream.read(keyed=True) == [ + {'header2': 'value2', 'header3': 'value3'}, + ] + + # Skip/pick columns def test_stream_skip_columns():