feat(connectors): let a Hub connector declare a source part by file pattern

packages/timenet-connectors/src/timenet_connectors/bases/huggingface.py

@@ -17,16 +17,16 @@
 
 from timenet.connectors import BaseConnector
 from timenet.dataset import TimeFDataset
-from timenet.errors import TimeFValidationError
-from timenet_connectors.sources.huggingface_hub import download_hub_files
+from timenet.errors import TimeFValidationError, TimeNetDownloadError
+from timenet_connectors.sources.huggingface_hub import hub_snapshot
 
 
 _COMMIT_SHA = re.compile(r"[0-9a-f]{40}")
 
 
 @dataclass(frozen=True)
 class HubFileRef:
-    """One local source file and the part of the dataset it holds."""
+    """One local source file, or the directory of a pattern's files, and the part of the dataset it holds."""
 
     path: Path
     part: str
@@ -46,7 +46,9 @@ class BaseHuggingFaceConnector(BaseConnector[HubFileRef], ABC):
     SOURCE_REVISION: ClassVar[str]
     """An immutable commit: a mutable branch can silently change rebuilt datasets."""
     SOURCE_FILES: ClassVar[Mapping[str, str]]
-    """Each part's repository-relative filename, in the order ``download()`` returns them."""
+    """Each part's repository-relative filename, in the order ``download()`` returns them. A filename can end
+    in a glob pattern, for example ``data/*.parquet``. Then the path of the part is the directory of the
+    files that match the pattern."""
 
     def __init_subclass__(cls, **kwargs: Any) -> None:
         super().__init_subclass__(**kwargs)
@@ -97,8 +99,12 @@ async def download_async(self, cache_dir: Path) -> list[HubFileRef]:
         Raises:
             ImportError: If ``huggingface_hub``, declared in the connector's requirements, is missing.
             TimeNetDownloadError: If a declared file is absent from the revision or cannot be downloaded.
-        """  # ruff: ignore[docstring-extraneous-exception] - both are raised by download_hub_files
-        paths = await download_hub_files(
-            self.HF_REPO, list(self.SOURCE_FILES.values()), revision=self.SOURCE_REVISION, cache_dir=cache_dir
-        )
-        return [HubFileRef(path, part) for part, path in zip(self.SOURCE_FILES, paths, strict=True)]
+        """  # ruff: ignore[docstring-extraneous-exception] - also raised by hub_snapshot
+        root = await hub_snapshot(self.HF_REPO, self.SOURCE_REVISION, cache_dir, tuple(self.SOURCE_FILES.values()))
+        refs = []
+        for part, name in self.SOURCE_FILES.items():
+            path = root / name if "*" not in name else root / Path(name).parent
+            if not path.exists():
+                raise TimeNetDownloadError(f"{self.HF_REPO}@{self.SOURCE_REVISION} has no {name}")
+            refs.append(HubFileRef(path, part))
+        return refs

packages/timenet-connectors/src/timenet_connectors/sources/huggingface_hub.py

@@ -1,17 +1,16 @@
 """Helpers for downloading files from pinned Hugging Face dataset revisions.
 
-:func:`download_hub_files` downloads with ``huggingface_hub.snapshot_download``, limited to the
-requested files. The Hub client runs the downloads in its own thread pool. Its progress bars stay
-hidden: the bar that counts file bytes reaches the ambient :mod:`~timenet_connectors.download.progress`
-sink instead, so the build CLI shows the download like any other connector download.
-``huggingface_hub`` is declared in the connector's ``requirements.txt`` and imported lazily, which
-keeps ``--no-isolation`` usable. The Hub client reads ``HF_TOKEN`` from the environment, so gated
-datasets work with no extra wiring.
+:func:`hub_snapshot` downloads with ``huggingface_hub.snapshot_download``. It downloads only the files
+that match the requested patterns, for example a filename or ``data/*.parquet``. The Hub client runs the
+downloads in its own thread pool. Its progress bars stay hidden. The bar that counts file bytes reaches
+the ambient :mod:`~timenet_connectors.download.progress` sink instead. Thus the build CLI shows the
+download like any other connector download. ``huggingface_hub`` is declared in the connector's ``requirements.txt``
+and imported lazily, which keeps ``--no-isolation`` usable. The Hub client reads ``HF_TOKEN`` from the
+environment, so gated datasets work with no extra wiring.
 """
 
 import asyncio
 from collections.abc import Sequence
-from fnmatch import fnmatchcase
 from pathlib import Path
 import threading
 from typing import Any
@@ -30,21 +29,22 @@
 )
 
 
-async def download_hub_files(repo: str, filenames: Sequence[str], *, revision: str, cache_dir: Path) -> list[Path]:
-    """Download files from a dataset revision with one ``snapshot_download`` call.
+async def hub_snapshot(repo: str, revision: str, cache_dir: Path, patterns: Sequence[str]) -> Path:
+    """Download the files of a dataset revision that match the patterns, with one ``snapshot_download`` call.
 
     Args:
         repo: The Hub dataset id, such as ``"BASH-Lab/SensorCaps"``.
-        filenames: Repository-relative paths of the files to download.
         revision: The revision to read, normally a pinned commit.
         cache_dir: Directory where the connector caches Hub files.
+        patterns: Repository-relative filenames or glob patterns, for example ``"data/*.parquet"``.
 
     Returns:
-        The local paths, in the order of ``filenames``.
+        The local snapshot directory. It holds the matching files at their repository paths.
+        The caller must check that the files it needs are there.
 
     Raises:
         ImportError: If ``huggingface_hub``, declared in the connector's requirements, is missing.
-        TimeNetDownloadError: If the download fails, or a file is absent from the revision.
+        TimeNetDownloadError: If the download fails.
     """
     try:
         # Optional connector dependency.
@@ -83,42 +83,11 @@ def update(self, n: float | None = 1) -> bool | None:
                     repo_type="dataset",
                     revision=revision,
                     cache_dir=cache_dir,
-                    allow_patterns=list(filenames),
+                    allow_patterns=list(patterns),
                     tqdm_class=_SinkProgress,
                 )
             )
         )
     except Exception as exc:
         raise TimeNetDownloadError(f"Could not download {repo}@{revision}: {exc}") from exc
-    paths = [snapshot / name for name in filenames]
-    missing = [name for name, path in zip(filenames, paths, strict=True) if not path.is_file()]
-    if missing:
-        raise TimeNetDownloadError(f"{repo}@{revision} has no {', '.join(missing)}")
-    return paths
-
-
-async def hub_snapshot(repo: str, revision: str, cache_dir: Path, patterns: tuple[str, ...]) -> Path:
-    """Download files that match the patterns from a pinned dataset revision.
-
-    Returns:
-        The local snapshot directory.
-
-    Raises:
-        ImportError: If ``huggingface_hub``, declared in the connector's requirements, is missing.
-        TimeNetDownloadError: If the revision cannot be listed or no files match.
-    """
-    try:
-        # Optional connector dependency.
-        from huggingface_hub import list_repo_files  # ruff: ignore[import-outside-top-level]
-    except ImportError as exc:
-        raise ImportError(_MISSING_HUB_CLIENT) from exc
-
-    try:
-        names = await asyncio.to_thread(list_repo_files, repo, repo_type="dataset", revision=revision)
-    except Exception as exc:
-        raise TimeNetDownloadError(f"Could not list {repo}@{revision}: {exc}") from exc
-    selected = [name for name in names if any(fnmatchcase(name, pattern) for pattern in patterns)]
-    if not selected:
-        raise TimeNetDownloadError(f"{repo}@{revision} has no files matching {patterns}")
-    paths = await download_hub_files(repo, selected, revision=revision, cache_dir=cache_dir)
-    return paths[0].parents[len(Path(selected[0]).parts) - 1]
+    return snapshot

packages/timenet-connectors/tests/test_huggingface_hub.py

@@ -1,6 +1,7 @@
 """Tests for the pinned Hub base and its download helper, driven against a fake huggingface_hub."""
 
 import asyncio
+from fnmatch import fnmatchcase
 from pathlib import Path
 import sys
 from types import SimpleNamespace
@@ -11,7 +12,7 @@
 from timenet.errors import TimeFValidationError, TimeNetDownloadError
 from timenet_connectors.bases.huggingface import BaseHuggingFaceConnector, HubFileRef
 from timenet_connectors.download.progress import DownloadProgress, progress_sink
-from timenet_connectors.sources.huggingface_hub import download_hub_files, hub_snapshot
+from timenet_connectors.sources.huggingface_hub import hub_snapshot
 
 
 _SHA = "0" * 40
@@ -37,8 +38,9 @@ def update(self, n=1):
         return None
 
 
-def _fake_hub(monkeypatch, tmp_path, present=None, error=None, bars=None):
-    """Install a fake huggingface_hub whose snapshot_download writes the allowed files into a snapshot."""
+def _fake_hub(monkeypatch, tmp_path, files=("data/train.jsonl", "data/test.jsonl"), error=None, bars=None):
+    """Install a fake huggingface_hub. Its snapshot_download writes the revision's ``files`` that match the
+    allowed patterns into a snapshot, as the Hub does."""
     calls = []
 
     def snapshot_download(repo, **kwargs):
@@ -48,8 +50,8 @@ def snapshot_download(repo, **kwargs):
         if error is not None:
             raise error
         snapshot = tmp_path / "snapshot"
-        for name in kwargs["allow_patterns"]:
-            if present is None or name in present:
+        for name in files:
+            if any(fnmatchcase(name, pattern) for pattern in kwargs["allow_patterns"]):
                 path = snapshot / name
                 path.parent.mkdir(parents=True, exist_ok=True)
                 path.write_text(name)
@@ -111,11 +113,9 @@ def test_download_failure_names_the_pinned_revision(monkeypatch, tmp_path):
 
 
 def test_a_file_absent_from_the_revision_fails(monkeypatch, tmp_path):
-    _fake_hub(monkeypatch, tmp_path, present={"data/train.jsonl"})
+    _fake_hub(monkeypatch, tmp_path, files=("data/train.jsonl",))
     with pytest.raises(TimeNetDownloadError, match=rf"org/repo@{_SHA} has no data/test\.jsonl$"):
-        asyncio.run(
-            download_hub_files("org/repo", ["data/train.jsonl", "data/test.jsonl"], revision=_SHA, cache_dir=tmp_path)
-        )
+        asyncio.run(_Pinned().download_async(tmp_path / "cache"))
 
 
 def test_missing_hub_library_states_the_fix(monkeypatch, tmp_path):
@@ -166,17 +166,40 @@ def test_convert_checks_refs_before_the_subclass_runs():
     assert _Pinned.convert.__name__ == "convert"
 
 
-def test_snapshot_selects_patterns_and_keeps_the_snapshot_root(monkeypatch, tmp_path):
-    calls = _fake_hub(monkeypatch, tmp_path)
-    hub = sys.modules["huggingface_hub"]
-    monkeypatch.setattr(
-        hub, "list_repo_files", lambda *args, **kwargs: ["meta.csv", "data/train.parquet", "README.md"], raising=False
-    )
+def test_snapshot_downloads_the_matching_files_and_returns_the_snapshot_root(monkeypatch, tmp_path):
+    calls = _fake_hub(monkeypatch, tmp_path, files=("meta.csv", "data/train.parquet", "README.md"))
 
     # asyncio.run gives the helper a running event loop, as download_async does in a build.
     root = asyncio.run(hub_snapshot("org/repo", _SHA, tmp_path, ("data/*.parquet",)))
     assert root == tmp_path / "snapshot"
-    assert (root / "data/train.parquet").is_file()
-    assert calls[0][1]["allow_patterns"] == ["data/train.parquet"]
-    with pytest.raises(TimeNetDownloadError, match="no files matching"):
-        asyncio.run(hub_snapshot("org/repo", _SHA, tmp_path, ("missing/*",)))
+    assert sorted(path.relative_to(root).as_posix() for path in root.rglob("*") if path.is_file()) == [
+        "data/train.parquet"
+    ]
+    assert calls[0][1]["allow_patterns"] == ["data/*.parquet"]
+
+
+class _Patterned(BaseHuggingFaceConnector):
+    HF_REPO = "org/repo"
+    SOURCE_REVISION = _SHA
+    SOURCE_FILES: ClassVar = {"questions": "qa.csv", "series": "data/*.parquet"}
+
+    def convert(self, raw_refs):
+        return raw_refs
+
+
+def test_a_pattern_part_resolves_to_the_directory_of_its_files(monkeypatch, tmp_path):
+    calls = _fake_hub(monkeypatch, tmp_path, files=("qa.csv", "data/a.parquet", "data/b.parquet", "README.md"))
+
+    refs = asyncio.run(_Patterned().download_async(tmp_path / "cache"))
+
+    assert refs == [
+        HubFileRef(tmp_path / "snapshot" / "qa.csv", "questions"),
+        HubFileRef(tmp_path / "snapshot" / "data", "series"),
+    ]
+    assert calls[0][1]["allow_patterns"] == ["qa.csv", "data/*.parquet"]
+
+
+def test_a_declared_file_or_pattern_without_a_match_fails(monkeypatch, tmp_path):
+    _fake_hub(monkeypatch, tmp_path, files=("qa.csv", "README.md"))
+    with pytest.raises(TimeNetDownloadError, match=rf"org/repo@{_SHA} has no data/\*\.parquet$"):
+        asyncio.run(_Patterned().download_async(tmp_path / "cache"))