FIX: derive csv header from the union of row keys

pyrit/common/csv_helper.py

@@ -20,13 +20,25 @@ def write_csv(file: IO[Any], examples: list[dict[str, str]]) -> None:
     """
     Write a list of dictionaries to a CSV file.
 
+    The header is the union of the keys across all rows, in first-seen order, so a
+    row carrying a key the first row lacks is written instead of failing halfway
+    through the file. Rows missing a header key get an empty cell.
+
     Args:
         file: A file-like object opened for writing CSV data.
         examples (list[dict[str, str]]): List of dictionaries to write as CSV rows.
     """
     if not examples:
         return
 
-    writer = csv.DictWriter(file, fieldnames=examples[0].keys())
+    fieldnames: list[str] = []
+    seen: set[str] = set()
+    for example in examples:
+        for key in example:
+            if key not in seen:
+                seen.add(key)
+                fieldnames.append(key)
+
+    writer = csv.DictWriter(file, fieldnames=fieldnames)
     writer.writeheader()
     writer.writerows(examples)

tests/unit/common/test_csv_helper.py

@@ -57,3 +57,27 @@ def test_write_csv_writes_header_and_rows():
     lines = file.getvalue().strip().splitlines()
     assert lines[0] == "name,role"
     assert lines[1] == "alice,admin"
+
+
+def test_write_csv_covers_keys_only_later_rows_have():
+    output = StringIO()
+    examples = [
+        {"name": "Alice"},
+        {"name": "Bob", "note": "extra"},
+    ]
+    write_csv(output, examples)
+    output.seek(0)
+    result = read_csv(output)
+    assert result == [{"name": "Alice", "note": ""}, {"name": "Bob", "note": "extra"}]
+
+
+def test_write_csv_keeps_first_seen_column_order():
+    output = StringIO()
+    examples = [
+        {"a": "1", "b": "2"},
+        {"c": "3", "a": "4"},
+    ]
+    write_csv(output, examples)
+    output.seek(0)
+    lines = output.read().strip().splitlines()
+    assert lines[0] == "a,b,c"