Normalize remote dataset file types from URLs

pyrit/datasets/seed_datasets/remote/remote_dataset_loader.py

@@ -10,6 +10,7 @@
 from collections.abc import Callable
 from pathlib import Path
 from typing import Any, Literal, Optional, TextIO, cast
+from urllib.parse import urlparse
 
 import requests
 from datasets import DownloadMode, disable_progress_bars, load_dataset
@@ -76,6 +77,18 @@ def _validate_file_type(self, file_type: str) -> None:
             valid_types = ", ".join(FILE_TYPE_HANDLERS.keys())
             raise ValueError(f"Invalid file_type. Expected one of: {valid_types}.")
 
+    def _get_file_type(self, *, source: str) -> str:
+        """
+        Infer the source file type from a URL or local path.
+
+        Query strings and fragments are ignored for URLs, and the result is
+        normalized to lowercase so `.JSON` and `.json` are treated identically.
+        """
+        parsed = urlparse(source)
+        source_path = parsed.path if parsed.scheme else source
+        suffix = Path(source_path).suffix
+        return suffix.lstrip(".").lower()
+
     def _read_cache(self, *, cache_file: Path, file_type: str) -> list[dict[str, str]]:
         """
         Read data from cache.
@@ -188,7 +201,7 @@ def _fetch_from_url(
             ...     source_type='public_url'
             ... )
         """
-        file_type = source.split(".")[-1]
+        file_type = self._get_file_type(source=source)
         if file_type not in FILE_TYPE_HANDLERS:
             valid_types = ", ".join(FILE_TYPE_HANDLERS.keys())
             raise ValueError(f"Invalid file_type. Expected one of: {valid_types}.")

tests/unit/datasets/test_remote_dataset_loader.py

@@ -72,3 +72,35 @@ def test_write_cache_creates_directories(self, tmp_path):
         loader._write_cache(cache_file=cache_file, examples=data, file_type="json")
 
         assert cache_file.exists()
+
+    @patch.object(_RemoteDatasetLoader, "_fetch_from_public_url", return_value=[{"key": "value"}])
+    def test_fetch_from_url_supports_query_string_file_type(self, mock_fetch_from_public_url):
+        loader = ConcreteRemoteLoader()
+
+        result = loader._fetch_from_url(
+            source="https://example.com/data.json?download=1",
+            source_type="public_url",
+            cache=False,
+        )
+
+        assert result == [{"key": "value"}]
+        mock_fetch_from_public_url.assert_called_once_with(
+            source="https://example.com/data.json?download=1",
+            file_type="json",
+        )
+
+    @patch.object(_RemoteDatasetLoader, "_fetch_from_public_url", return_value=[{"key": "value"}])
+    def test_fetch_from_url_supports_uppercase_file_type(self, mock_fetch_from_public_url):
+        loader = ConcreteRemoteLoader()
+
+        result = loader._fetch_from_url(
+            source="https://example.com/data.JSON",
+            source_type="public_url",
+            cache=False,
+        )
+
+        assert result == [{"key": "value"}]
+        mock_fetch_from_public_url.assert_called_once_with(
+            source="https://example.com/data.JSON",
+            file_type="json",
+        )