diff --git a/forge/cli.py b/forge/cli.py index 53ad10f..b0f0487 100644 --- a/forge/cli.py +++ b/forge/cli.py @@ -207,6 +207,38 @@ def _resolve_dataset_path(path_str: str, demo: bool = False) -> Path: return path +def _registry_id_to_hf_url(path_str: str, demo: bool = False) -> str | None: + """Map a bare registry id to its ``hf://`` source URL, if it has one. + + A registry id like ``metaworld`` otherwise skips the ``--quick`` and + size-guard branches in ``inspect`` (which only recognise HF URLs), so the + whole dataset downloads with no metadata-only option and no size warning. + Rewriting the id to ``hf://`` up front lets those branches apply + uniformly. Non-HF sources (gcs/http/…) return ``None`` so the caller falls + back to ``_resolve_dataset_path``, which keeps their existing handling. + """ + if "/" in path_str or path_str.startswith((".", "/")): + return None + if Path(path_str).exists(): + return None + try: + from forge.registry import DatasetRegistry + except ImportError: + return None + try: + entry = DatasetRegistry.get(path_str) + source = DatasetRegistry.get_source(path_str, demo=demo) + except Exception: + # Not a registry id (or no usable source) — let the normal path report it. + return None + if source.type != "hf_hub": + return None + console.print( + f"[cyan]Resolved from registry:[/cyan] {entry.name} ({entry.format})" + ) + return f"hf://{source.uri}" + + def _quick_inspect_hub(path: str, output: str = "text") -> None: """Quick inspect a HuggingFace Hub dataset without downloading. @@ -380,6 +412,13 @@ def inspect_cmd( from forge.core.exceptions import ForgeError from forge.hub import is_hf_url + # Rewrite a bare registry id (e.g. `metaworld`) to its hf:// source so the + # --quick and size-guard branches below apply to registry ids too, not just + # explicit HF URLs. Non-HF-backed ids are left as-is for _resolve_dataset_path. + hf_from_registry = _registry_id_to_hf_url(path) + if hf_from_registry is not None: + path = hf_from_registry + # Quick inspect for Hub datasets (metadata only, no download) if quick and is_hf_url(path): _quick_inspect_hub(path, output) diff --git a/tests/test_cli.py b/tests/test_cli.py index 15817a5..030bbc2 100644 --- a/tests/test_cli.py +++ b/tests/test_cli.py @@ -171,3 +171,38 @@ def test_lerobot_layout(self, tmp_path): ) assert result.exit_code == 0, result.stdout assert "myorg/stack_lego" in result.stdout + + +class TestRegistryIdToHfUrl: + """Registry ids should be rewritten to hf:// so --quick / size-guard apply.""" + + def test_hf_backed_id_resolves_to_hf_url(self): + from forge.cli import _registry_id_to_hf_url + + # metaworld's source is hf_hub: lerobot/metaworld_mt50 + assert _registry_id_to_hf_url("metaworld") == "hf://lerobot/metaworld_mt50" + + def test_prefers_hf_over_other_sources(self): + from forge.cli import _registry_id_to_hf_url + + # droid has both a gcs and an hf_hub source; get_source prefers hf_hub. + assert _registry_id_to_hf_url("droid") == "hf://cadene/droid" + + def test_unknown_id_returns_none(self): + from forge.cli import _registry_id_to_hf_url + + assert _registry_id_to_hf_url("definitely_not_a_dataset") is None + + def test_repo_id_and_paths_are_ignored(self): + from forge.cli import _registry_id_to_hf_url + + assert _registry_id_to_hf_url("lerobot/pusht") is None + assert _registry_id_to_hf_url("./local_dir") is None + + def test_local_dir_shadows_registry_id(self, tmp_path, monkeypatch): + from forge.cli import _registry_id_to_hf_url + + # A local directory named like a registry id must not be rewritten. + monkeypatch.chdir(tmp_path) + (tmp_path / "metaworld").mkdir() + assert _registry_id_to_hf_url("metaworld") is None