Accept Iterable[DatasetEntry] in transformations and exporters

Co-Authored-By: Claude Opus 4.7 <noreply@anthropic.com>

Author: J03D03

examples/04_enrich_with_commit_data.py

@@ -8,7 +8,6 @@
 setup_logging("enrich_with_commit_data")
 
 if __name__ == "__main__":
-    entries = list(vfc_datasets.DevignDataset()) # Devign contains only 2 projects (FFmpeg, QEMU).
-
-    entries = transformations.add_commit_information_local(entries)
+    # Devign contains only 2 projects (FFmpeg, QEMU).
+    entries = transformations.add_commit_information_local(vfc_datasets.DevignDataset())
     log_dataset_stats(entries)

examples/05_save_and_export.py

@@ -9,7 +9,7 @@
 setup_logging("save_and_export")
 
 if __name__ == "__main__":
-    entries = list(vfc_datasets.DevignDataset())
+    entries = vfc_datasets.DevignDataset()
 
     output_dir = Path(".data/exports")
     save_entries(entries, output_dir / "devign.jsonl")

examples/06_create_splits.py

@@ -9,10 +9,8 @@
 setup_logging("create_splits")
 
 if __name__ == "__main__":
-    entries = list(vfc_datasets.DevignDataset())
-
     create_random_split(
-        entries,
+        vfc_datasets.DevignDataset(),
         name="devign",
         output_path=Path(".data/splits"),
         seed=42,

examples/07_filter_by_language.py

@@ -10,8 +10,8 @@
 logger = logging.getLogger(__name__)
 
 if __name__ == "__main__":
-    entries = list(vfc_datasets.BigVulDataset())
+    dataset = vfc_datasets.BigVulDataset()
 
-    logger.info("Entries before filtering: %d", len(entries))
-    entries = transformations.filter_by_extension(entries, {"c", "h", "cpp", "hpp", "cc"})
+    logger.info("Entries before filtering: %d", len(dataset))
+    entries = transformations.filter_by_extension(dataset, {"c", "h", "cpp", "hpp", "cc"})
     logger.info("Entries after filtering: %d", len(entries))

src/vfc_datasets/transformations/enrichment/add_commit_data_api.py

@@ -2,6 +2,7 @@
 
 import asyncio
 import logging
+from collections.abc import Iterable
 from typing import Any
 
 from tqdm.asyncio import tqdm
@@ -75,8 +76,9 @@ async def _enrich_entries_async(entries: list[DatasetEntry]) -> tuple[int, int]:
     return success, len(entries) - success
 
 
-def add_commit_information_api(entries: list[DatasetEntry]) -> list[DatasetEntry]:
+def add_commit_information_api(entries: Iterable[DatasetEntry]) -> list[DatasetEntry]:
     """Enrich entries with commit data from the GitHub API and return the modified list."""
+    entries = list(entries)
     entries_to_process = [e for e in entries if needs_enrichment(e)]
 
     if not entries_to_process:

src/vfc_datasets/transformations/enrichment/add_commit_data_local.py

@@ -1,4 +1,5 @@
 import logging
+from collections.abc import Iterable
 from datetime import UTC, datetime
 
 from git import Repo
@@ -53,13 +54,13 @@ def _process_commit_batch(args: tuple[str, list[str], int]) -> dict[str, CommitD
     return results
 
 
-def add_commit_information_local(entries: list[DatasetEntry]) -> list[DatasetEntry]:
+def add_commit_information_local(entries: Iterable[DatasetEntry]) -> list[DatasetEntry]:
     """Enrich DatasetEntry objects with commit information from local git repos."""
     logger.info("Add commit information [LOCAL]")
     logger.info("Max diff size limit: %dK chars", MAX_DIFF_SIZE // 1000)
 
     return process_commits_in_batches(
-        entries,
+        list(entries),
         filter_fn=needs_enrichment,
         batch_fn=_process_commit_batch,
         apply_fn=apply_commit_data,

src/vfc_datasets/transformations/enrichment/add_no_comment.py

@@ -1,6 +1,7 @@
 import logging
 import os
 import tempfile
+from collections.abc import Iterable
 from pathlib import PurePosixPath
 from typing import Any
 
@@ -290,18 +291,19 @@ def _apply_diff(entry: DatasetEntry, diff: str) -> None:
 
 
 def strip_diff_comments(
-    entries: list[DatasetEntry], *, include_unsupported: bool = True
+    entries: Iterable[DatasetEntry], *, include_unsupported: bool = True
 ) -> list[DatasetEntry]:
     """Strip comments from commit diffs in-place using tree-sitter.
 
     Args:
-        entries: List of dataset entries to process
+        entries: Dataset entries to process
         include_unsupported: If True (default), include original diff for unsupported files.
                             If False, skip unsupported files entirely.
     """
     logger.info("Strip comments from commit diffs [LOCAL]")
     logger.info("Max diff size: %dK chars", MAX_DIFF_SIZE // 1000)
 
+    entries = list(entries)
     needs_processing = [e for e in entries if e.commit_diff]
     skipped = sum(
         1

src/vfc_datasets/transformations/enrichment/commit_id_enrichment.py

@@ -1,6 +1,7 @@
 import asyncio
 import logging
 from collections import defaultdict
+from collections.abc import Iterable
 from concurrent.futures import ProcessPoolExecutor, as_completed
 
 from tqdm.asyncio import tqdm as async_tqdm
@@ -114,9 +115,9 @@ async def _extend_commit_id_api_async(
     return entry, entry.commit_id, False
 
 
-def extend_commit_ids_api(entries: list[DatasetEntry]) -> list[DatasetEntry]:
+def extend_commit_ids_api(entries: Iterable[DatasetEntry]) -> list[DatasetEntry]:
     """Extend commit IDs using the GitHub API (async) and return the (possibly) modified list of entries."""
-    # Filter entries that need extension
+    entries = list(entries)
     entries_to_process = [
         entry for entry in entries if entry.commit_id and len(entry.commit_id) < 40
     ]
@@ -139,8 +140,9 @@ def extend_commit_ids_api(entries: list[DatasetEntry]) -> list[DatasetEntry]:
     return entries
 
 
-def extend_commit_ids_local(entries: list[DatasetEntry]) -> list[DatasetEntry]:
+def extend_commit_ids_local(entries: Iterable[DatasetEntry]) -> list[DatasetEntry]:
     """Extend commit IDs in-place by cloning repositories locally and return the list of modified entries."""
+    entries = list(entries)
     entries_to_process = [
         entry for entry in entries if entry.commit_id and len(entry.commit_id) < 40
     ]

src/vfc_datasets/transformations/enrichment/project_urls/update_project_urls.py

@@ -1,11 +1,12 @@
 """Transformation functions for updating and filtering project URLs.
 
-These functions operate on a list of DatasetEntry objects to:
+These functions operate on dataset entries to:
 - Remove entries with unreachable project URLs
 - Update entries with moved/renamed project URLs
 """
 
 import logging
+from collections.abc import Iterable
 
 from vfc_datasets.dataset_entry import DatasetEntry
 
@@ -17,16 +18,17 @@
 logger = logging.getLogger(__name__)
 
 
-def filter_unreachable_project_urls(entries: list[DatasetEntry]) -> list[DatasetEntry]:
+def filter_unreachable_project_urls(entries: Iterable[DatasetEntry]) -> list[DatasetEntry]:
     """Remove entries with unreachable project_url's"""
     logger.info("REMOVE unreachable project_urls")
     unreachable_urls = get_unreachable_urls()
     return [e for e in entries if e.project_url not in unreachable_urls]
 
 
-def update_project_urls_inplace(entries: list[DatasetEntry]) -> list[DatasetEntry]:
+def update_project_urls_inplace(entries: Iterable[DatasetEntry]) -> list[DatasetEntry]:
     """Update moved or fixed project_url's in-place and return the modified entries."""
     logger.info("UPDATE project_urls")
+    entries = list(entries)
     changed_urls = 0
     moved_urls = get_moved_urls()
     for entry in entries:

src/vfc_datasets/transformations/filters/collapse_to_commit_level.py

@@ -1,5 +1,6 @@
 import logging
 from collections import defaultdict
+from collections.abc import Iterable
 
 from vfc_datasets.dataset_entry import DatasetEntry
 
@@ -9,7 +10,7 @@
 
 
 def collapse_to_commit_level(
-    entries: list[DatasetEntry],
+    entries: Iterable[DatasetEntry],
     *,
     include_benign_only: bool = True,
 ) -> list[DatasetEntry]:
@@ -23,13 +24,13 @@ def collapse_to_commit_level(
             "functions. This could lead to mislabeling VFCs as non-VFC."
         )
 
-    if not entries:
-        return []
-
     groups: dict[tuple[str, str], list[DatasetEntry]] = defaultdict(list)
     for entry in entries:
         groups[(entry.project_url, entry.commit_id)].append(entry)
 
+    if not groups:
+        return []
+
     result = []
     dropped = 0
     for group_entries in groups.values():