diff --git a/tests/data/test_dataset.py b/tests/data/test_dataset.py index aae1f753..1049c2d2 100755 --- a/tests/data/test_dataset.py +++ b/tests/data/test_dataset.py @@ -356,3 +356,22 @@ def test_full_trajectory_mode_minimum_steps(tmp_path, start_output_steps_at_t): assert ( ctrl_data["output_time_grid"][offset - 2] == exp_data["input_time_grid"][0] ), "Expected input time to be time prior to first target output time" + + +def test_join_path_keeps_forward_slashes_for_remote_uris(monkeypatch): + """Remote URIs must not get Windows separators (see issue #91).""" + import ntpath + + from the_well.data import datasets + + monkeypatch.setattr(datasets.os, "path", ntpath) + + remote = datasets._join_path( + "hf://datasets/polymathic-ai/", "turbulent_radiative_layer_2D", "data", "train" + ) + assert ( + remote == "hf://datasets/polymathic-ai/turbulent_radiative_layer_2D/data/train" + ) + + local = datasets._join_path("C:\\well", "turbulent_radiative_layer_2D", "data") + assert local == "C:\\well\\turbulent_radiative_layer_2D\\data" diff --git a/the_well/data/datasets.py b/the_well/data/datasets.py index 83c5115c..d78466f0 100755 --- a/the_well/data/datasets.py +++ b/the_well/data/datasets.py @@ -1,5 +1,6 @@ import itertools import os +import posixpath import warnings from dataclasses import dataclass from enum import Enum @@ -108,6 +109,22 @@ class TrajectoryMetadata: time_stride: int | List[int] +def _is_remote_path(path) -> bool: + """Whether `path` is a URI handled by fsspec (e.g. ``hf://``, ``s3://``).""" + return "://" in str(path) + + +def _join_path(base, *parts: str) -> str: + """Join path components, keeping forward slashes for remote URIs. + + `os.path.join` inserts backslashes on Windows, which breaks fsspec URIs + such as ``hf://datasets/...``. Remote paths are always POSIX-style. + """ + if _is_remote_path(base): + return posixpath.join(str(base), *parts) + return os.path.join(base, *parts) + + class WellDataset(Dataset): """ Generic dataset for any Well data. Returns data in B x T x H [x W [x D]] x C format. @@ -228,24 +245,24 @@ def __init__( if normalization_path is None: normalization_path = "stats.yaml" if well_split_name is not None: - self.data_path = os.path.join(path, "data", well_split_name) + self.data_path = _join_path(path, "data", well_split_name) else: assert is_dataset_in_the_well( well_dataset_name ), f"Dataset name {well_dataset_name} not in the expected list {WELL_DATASETS}." - self.data_path = os.path.join( + self.data_path = _join_path( well_base_path, well_dataset_name, "data", well_split_name ) - trunk_path = os.path.join(well_base_path, well_dataset_name) + trunk_path = _join_path(well_base_path, well_dataset_name) if normalization_path is None: - normalization_path = os.path.join( + normalization_path = _join_path( well_base_path, well_dataset_name, "stats.yaml" ) - if os.path.isabs(normalization_path): + if os.path.isabs(normalization_path) or _is_remote_path(normalization_path): self.normalization_path = normalization_path else: - self.normalization_path = os.path.join( + self.normalization_path = _join_path( trunk_path, str(normalization_path).removeprefix(str(trunk_path)).lstrip("./\\"), )