Keep forward slashes when joining remote dataset paths

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"

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("./\\"),
             )