diff --git a/src/hdx/utilities/saver.py b/src/hdx/utilities/saver.py index 2c4064b..fd20984 100644 --- a/src/hdx/utilities/saver.py +++ b/src/hdx/utilities/saver.py @@ -330,9 +330,15 @@ def write_rows(newrows, has_header, headers) -> None: if isinstance(row, dict): has_header = True row = row_function(row) - while row is None: - row = next(rows) - row = row_function(row) + try: + while row is None: + row = next(rows) + row = row_function(row) + except StopIteration: + if not no_empty and headers: + newrows = [headers] + write_rows(newrows, None, None) + return newrows if columns: row = {k: row[k] for k in columns if k in row} newrows.append(row) diff --git a/tests/hdx/utilities/test_saver.py b/tests/hdx/utilities/test_saver.py index 840fc7c..6477e52 100755 --- a/tests/hdx/utilities/test_saver.py +++ b/tests/hdx/utilities/test_saver.py @@ -296,7 +296,7 @@ def test_save_hxlated_output(self, tmp_path, saverfolder, json_csv_configuration filename = "out8.json" assert_files_same(saverfolder / filename, output_dir / filename) - def test_save_iterable(self): + def test_save_iterable_lists(self): list_of_tuples = [(1, 2, 3, "a"), (4, 5, 6, "b"), (7, 8, 9, "c")] list_of_lists = [[1, 2, 3, "a"], [4, 5, 6, "b"], [7, 8, 9, "c"]] list_of_dicts = [ @@ -457,6 +457,15 @@ def row_func(row): ["7", "8", "9", "c"], ] + def test_save_iterable_dicts(self): + with temp_dir( + "TestSaveIterable", + delete_on_success=True, + delete_on_failure=False, + ) as tempdir: + filename = "test_save_iterable_to_csv.csv" + filepath = tempdir / filename + list_of_dicts = [ {"h1": 1, "h2": 2, "h3": 3, "h4": "a"}, {"h1": 4, "h2": 5, "h3": 6, "h4": "b"}, @@ -554,3 +563,47 @@ def row_func(row): newll = read_list_from_csv(str(filepath)) remove(filepath) assert newll == [["h1", "h3", "h4"]] + + def process_row(row): + return None + + rows = save_iterable( + filepath, + list_of_dicts, + headers=["h1", "h2", "h3", "h4"], + row_function=process_row, + ) + assert rows == [] + + rows = save_iterable( + filepath, + list_of_dicts, + headers=["h1", "h2", "h3", "h4"], + row_function=process_row, + no_empty=False, + ) + assert rows == [["h1", "h2", "h3", "h4"]] + newll = read_list_from_csv(filepath) + remove(filepath) + assert newll == [ + ["h1", "h2", "h3", "h4"], + ] + + def process_row(row): + if 5 in row.values(): + return None + return row + + save_iterable( + filepath, + list_of_dicts, + headers=["h1", "h2", "h3", "h4"], + row_function=process_row, + ) + newll = read_list_from_csv(filepath) + remove(filepath) + assert newll == [ + ["h1", "h2", "h3", "h4"], + ["1", "2", "4", "a"], + ["7", "8", "10", "c"], + ]