Merge pull request #98 from fish-pace/copilot/fix-open-dataset-respecting-method

open_dataset: remove open_dataset_kwargs, add silent param, integer index, datatree preset, and fix show_variables auto-resolution

Author: eeholmes

src/point_collocation/core/_open_method.py

@@ -63,7 +63,7 @@
 )
 
 _VALID_XARRAY_OPEN: frozenset[str] = frozenset({"dataset", "datatree"})
-_VALID_PRESETS: frozenset[str] = frozenset({"dataset", "datatree-merge", "auto"})
+_VALID_PRESETS: frozenset[str] = frozenset({"dataset", "datatree", "datatree-merge", "auto"})
 
 # Default open kwargs applied to both xr.open_dataset and xr.open_datatree
 # unless explicitly overridden by the user.
@@ -163,7 +163,8 @@ def _expand_preset(preset: str) -> dict:
     Parameters
     ----------
     preset:
-        One of ``"dataset"``, ``"datatree-merge"``, or ``"auto"``.
+        One of ``"dataset"``, ``"datatree"``, ``"datatree-merge"``, or
+        ``"auto"``.
 
     Returns
     -------
@@ -179,6 +180,17 @@ def _expand_preset(preset: str) -> dict:
         return {
             "xarray_open": "dataset",
             "open_kwargs": {},
+            "merge": None,
+            "coords": "auto",
+            "set_coords": True,
+            "dim_renames": None,
+            "auto_align_phony_dims": None,
+        }
+    if preset == "datatree":
+        return {
+            "xarray_open": "datatree",
+            "open_kwargs": {},
+            "merge": None,
             "coords": "auto",
             "set_coords": True,
             "dim_renames": None,
@@ -252,11 +264,13 @@ def _validate_and_fill_spec(spec: dict) -> dict:
         )
 
     if xarray_open == "datatree":
-        result.setdefault("merge", "all")
-        result.setdefault("merge_kwargs", {})
-    elif xarray_open == "dataset" and "merge" in result:
-        # Allow merge_kwargs alongside merge for the dataset path too.
-        result.setdefault("merge_kwargs", {})
+        result.setdefault("merge", None)
+        if result.get("merge") is not None:
+            result.setdefault("merge_kwargs", {})
+    elif xarray_open == "dataset":
+        result.setdefault("merge", None)
+        if result.get("merge") is not None:
+            result.setdefault("merge_kwargs", {})
 
     return result
 

src/point_collocation/core/engine.py

@@ -88,6 +88,7 @@ def matchup(
 
         * ``"dataset"`` — open with ``xarray.open_dataset`` (fast path for
           typical flat NetCDF files).
+        * ``"datatree"`` — open as a raw DataTree without merging groups.
         * ``"datatree-merge"`` — open as DataTree and merge all groups into
           a flat Dataset (for grouped/HDF5-ish files).
         * ``"auto"`` *(default)* — try the fast ``"dataset"`` path first; if
@@ -101,7 +102,7 @@ def matchup(
             open_method = {
                 "xarray_open":           "dataset" | "datatree",
                 "open_kwargs":           {},
-                "merge":                 "all" | "root" | ["/path/a"],
+                "merge":                 None | "all" | "root" | ["/path/a"],
                 "merge_kwargs":          {},
                 "coords":                "auto" | ["Lat", "Lon"] | {"lat": "...", "lon": "..."},
                 "set_coords":            True,

src/point_collocation/core/plan.py

@@ -188,39 +188,72 @@ def __getitem__(self, idx: int | slice) -> "Plan | Any":
 
     def open_dataset(
         self,
-        result: Any,
-        open_method: str | dict | None = None,
-        open_dataset_kwargs: dict[str, Any] | None = None,
-    ) -> "xr.Dataset":
-        """Open a single granule result as an :class:`xarray.Dataset`.
+        result: "int | Any",
+        open_method: "str | dict | None" = None,
+        *,
+        silent: bool = False,
+    ) -> "Any":
+        """Open a single granule result as an :class:`xarray.Dataset` or DataTree.
 
         Parameters
         ----------
         result:
-            A single earthaccess result object, typically obtained via
-            ``plan[n]``.
+            An integer index into ``plan.results`` (e.g. ``0``), or a
+            single earthaccess result object obtained via ``plan[n]``.
+            Using an integer is preferred: ``plan.open_dataset(0)`` is
+            equivalent to ``plan.open_dataset(plan[0])``.
         open_method:
             How to open the granule.  Accepts the same string presets or
             dict spec as :func:`~point_collocation.matchup`.  Defaults to
             ``"auto"`` (try dataset first, fall back to datatree merge).
-        open_dataset_kwargs:
-            Keyword arguments forwarded to the xarray open function.
-            ``chunks`` defaults to ``{}`` (lazy/dask loading).  ``engine``
-            defaults to ``"h5netcdf"`` when not specified.
+
+            **String presets:**
+
+            * ``"dataset"`` — open with ``xarray.open_dataset`` (flat NetCDF).
+            * ``"datatree"`` — open as a DataTree with all groups; returns the
+              raw :class:`xarray.DataTree` (or ``datatree.DataTree``) without
+              merging groups.  Equivalent to ``xarray.open_datatree(f)``.
+            * ``"datatree-merge"`` — open as DataTree and merge all groups into
+              a flat Dataset.
+            * ``"auto"`` *(default)* — probe the file first; if lat/lon can be
+              detected via ``xr.open_dataset``, use that; otherwise fall back to
+              ``"datatree-merge"``.  The printed spec shows the **resolved** mode.
+
+            Pass open-function kwargs via the ``"open_kwargs"`` key of a
+            dict spec, e.g.
+            ``open_method={"open_kwargs": {"engine": "netcdf4"}}``.
+        silent:
+            When ``False`` (default), print the effective open_method spec
+            actually used (after normalization and auto-resolution).
+            Set to ``True`` to suppress this output.
 
         Returns
         -------
-        xarray.Dataset
-            The caller is responsible for closing the dataset when finished
-            (e.g. ``ds.close()`` or ``with plan.open_dataset(...) as ds:``).
+        xarray.Dataset or xarray.DataTree
+            A flat :class:`xarray.Dataset` for all modes except
+            ``open_method="datatree"`` (or a dict spec with
+            ``xarray_open="datatree"`` and ``merge=None``), which returns the
+            raw DataTree.
+            The caller is responsible for closing the returned object when
+            finished (e.g. ``ds.close()``).
         """
+        if isinstance(result, int):
+            n = len(self.results)
+            if result < 0 or result >= n:
+                raise IndexError(
+                    f"result index {result} is out of range for a plan with {n} result(s). "
+                    f"Valid indices are 0 to {n - 1}."
+                )
+            result = self.results[result]
+
         from point_collocation.core._open_method import (
             _apply_coords,
             _build_effective_open_kwargs,
             _merge_datatree_with_spec,
             _normalize_open_method,
             _open_and_merge_dataset_groups,
             _open_datatree_fn,
+            _resolve_auto_spec,
         )
 
         try:
@@ -234,7 +267,10 @@ def open_dataset(
         import xarray as xr
 
         effective_open_method = "auto" if open_method is None else open_method
-        spec = _normalize_open_method(effective_open_method, open_dataset_kwargs)
+        spec = _normalize_open_method(effective_open_method)
+
+        xarray_open = spec.get("xarray_open", "dataset")
+        effective_kwargs = _build_effective_open_kwargs(spec.get("open_kwargs", {}))
 
         file_objs = earthaccess.open([result], pqdm_kwargs={"disable": True})
         if len(file_objs) != 1:
@@ -243,10 +279,28 @@ def open_dataset(
             )
         file_obj = file_objs[0]
 
-        xarray_open = spec.get("xarray_open", "dataset")
-        effective_kwargs = _build_effective_open_kwargs(spec.get("open_kwargs", {}))
+        # For "auto" mode, probe the file first so that the printed spec shows
+        # the actual resolved mode (e.g. "dataset" or "datatree"), not "auto".
+        if xarray_open == "auto":
+            try:
+                spec = _resolve_auto_spec(file_obj, spec)
+                xarray_open = spec["xarray_open"]
+                effective_kwargs = _build_effective_open_kwargs(spec.get("open_kwargs", {}))
+            except ValueError:
+                xarray_open = "dataset"
+                spec = {**spec, "xarray_open": "dataset", "merge": None}
+                effective_kwargs = _build_effective_open_kwargs(spec.get("open_kwargs", {}))
+
+        if not silent:
+            display_spec = {**spec, "open_kwargs": effective_kwargs}
+            print(f"open_method: {display_spec!r}")
 
         if xarray_open == "datatree":
+            merge = spec.get("merge")
+            if merge is None:
+                # Return the raw DataTree without merging — like open_datatree(f).
+                return _open_datatree_fn(file_obj, effective_kwargs)
+            # merge is "all", "root", or a list: merge groups into a flat Dataset.
             dt = _open_datatree_fn(file_obj, effective_kwargs)
             try:
                 ds = _merge_datatree_with_spec(dt, spec)
@@ -259,11 +313,7 @@ def open_dataset(
                 pass  # coords not found; return merged dataset as-is
             return ds
 
-        if xarray_open in ("dataset", "auto"):
-            # For the dataset and auto paths, return an open dataset whose
-            # lifecycle is managed by the caller.  Auto tries the fast path
-            # only; if the caller needs datatree fallback they should use
-            # open_method="datatree-merge" explicitly.
+        if xarray_open == "dataset":
             merge = spec.get("merge")
             if merge is not None:
                 # Dataset-based group merge: open each group and merge.
@@ -283,8 +333,9 @@ def open_dataset(
     def open_mfdataset(
         self,
         results: "list[Any] | Plan",
-        open_method: str | dict | None = None,
-        open_dataset_kwargs: dict[str, Any] | None = None,
+        open_method: "str | dict | None" = None,
+        *,
+        silent: bool = False,
     ) -> "xr.Dataset":
         """Open multiple granule results as a single :class:`xarray.Dataset`.
 
@@ -300,10 +351,13 @@ def open_mfdataset(
             ``"datatree-merge"`` opens each granule as a DataTree, merges
             its groups into a flat dataset, then concatenates all granules
             along a new ``granule`` dimension.  Defaults to ``"auto"``.
-        open_dataset_kwargs:
-            Keyword arguments forwarded to the xarray open function.
-            ``chunks`` defaults to ``{}`` (lazy/dask loading).  ``engine``
-            defaults to ``"h5netcdf"`` when not specified.
+            Pass open-function kwargs via the ``"open_kwargs"`` key of a
+            dict spec, e.g.
+            ``open_method={"open_kwargs": {"engine": "netcdf4"}}``.
+        silent:
+            When ``False`` (default), print the effective open_method spec
+            actually used (after normalization and defaults are applied).
+            Set to ``True`` to suppress this output.
 
         Returns
         -------
@@ -329,17 +383,21 @@ def open_mfdataset(
         import xarray as xr
 
         effective_open_method = "auto" if open_method is None else open_method
-        spec = _normalize_open_method(effective_open_method, open_dataset_kwargs)
+        spec = _normalize_open_method(effective_open_method)
+
+        xarray_open = spec.get("xarray_open", "dataset")
+        effective_kwargs = _build_effective_open_kwargs(spec.get("open_kwargs", {}))
+
+        if not silent:
+            display_spec = {**spec, "open_kwargs": effective_kwargs}
+            print(f"open_method: {display_spec!r}")
 
         result_list = results.results if isinstance(results, Plan) else list(results)
         file_objs = earthaccess.open(result_list, pqdm_kwargs={"disable": True})
 
-        xarray_open = spec.get("xarray_open", "dataset")
-
         if xarray_open == "datatree":
             # Open each granule as a DataTree, merge its groups, then
             # concatenate all granule datasets along a new "granule" dim.
-            effective_kwargs = _build_effective_open_kwargs(spec.get("open_kwargs", {}))
             merged_datasets: list[xr.Dataset] = []
             for file_obj in file_objs:
                 dt = _open_datatree_fn(file_obj, effective_kwargs)
@@ -357,7 +415,6 @@ def open_mfdataset(
             # separate datasets and merge them, then concatenate all granules
             # along a new "granule" dimension.
             # Without merge, use xr.open_mfdataset for simplicity.
-            effective_kwargs = _build_effective_open_kwargs(spec.get("open_kwargs", {}))
             merge = spec.get("merge")
             if merge is not None:
                 merged_datasets = []
@@ -380,8 +437,7 @@ def open_mfdataset(
 
     def show_variables(
         self,
-        open_method: str | dict | None = None,
-        open_dataset_kwargs: dict[str, Any] | None = None,
+        open_method: "str | dict | None" = None,
     ) -> None:
         """Print the first granule's groups, dimensions, and variables.
 
@@ -395,19 +451,16 @@ def show_variables(
         the merged flat dataset.
 
         To obtain the full xarray representation programmatically use
-        ``plan.open_dataset(plan[0])`` instead.
+        ``plan.open_dataset(0)`` instead.
 
         Parameters
         ----------
         open_method:
             How to open the granule.  Accepts the same string presets or
             dict spec as :func:`~point_collocation.matchup`.  Defaults to
-            ``"auto"``.
-        open_dataset_kwargs:
-            Keyword arguments forwarded to the xarray open function when
-            the h5py fast path is unavailable.
-            ``chunks`` defaults to ``{}`` (lazy/dask loading).  ``engine``
-            defaults to ``"h5netcdf"`` when not specified.
+            ``"auto"``.  Pass open-function kwargs via the ``"open_kwargs"``
+            key of a dict spec, e.g.
+            ``open_method={"open_kwargs": {"engine": "netcdf4"}}``.
 
         Raises
         ------
@@ -430,7 +483,7 @@ def show_variables(
             raise ValueError("No granules in plan — cannot show variables.")
 
         effective_open_method = "auto" if open_method is None else open_method
-        spec = _normalize_open_method(effective_open_method, open_dataset_kwargs)
+        spec = _normalize_open_method(effective_open_method)
         xarray_open = spec.get("xarray_open", "dataset")
 
         try:
@@ -452,6 +505,18 @@ def show_variables(
 
         effective_kwargs = _build_effective_open_kwargs(spec.get("open_kwargs", {}))
 
+        # For "auto" mode, probe the file first so the printed spec shows the
+        # actual resolved mode (e.g. "dataset"), not "auto".
+        if xarray_open == "auto":
+            try:
+                spec = _resolve_auto_spec(file_obj, spec)
+                xarray_open = spec["xarray_open"]
+                effective_kwargs = _build_effective_open_kwargs(spec.get("open_kwargs", {}))
+            except ValueError:
+                xarray_open = "dataset"
+                spec = {**spec, "xarray_open": "dataset", "merge": None}
+                effective_kwargs = _build_effective_open_kwargs(spec.get("open_kwargs", {}))
+
         # Print the spec summary.
         display_spec = {**spec, "open_kwargs": effective_kwargs}
         print(f"open_method: {display_spec!r}")