feat: open_dataset respects open_method, removes open_dataset_kwargs, adds silent param and integer index support

Co-authored-by: eeholmes <[email protected]>

Author: Copilot

src/point_collocation/core/plan.py

@@ -188,32 +188,47 @@ 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,
+        result: "int | Any",
+        open_method: "str | dict | None" = None,
+        *,
+        silent: bool = False,
     ) -> "xr.Dataset":
         """Open a single granule result as an :class:`xarray.Dataset`.
 
         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.
+            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
         -------
         xarray.Dataset
             The caller is responsible for closing the dataset when finished
             (e.g. ``ds.close()`` or ``with plan.open_dataset(...) as ds:``).
         """
+        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,
@@ -234,7 +249,14 @@ 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", {}))
+
+        if not silent:
+            display_spec = {**spec, "open_kwargs": effective_kwargs}
+            print(f"open_method: {display_spec!r}")
 
         file_objs = earthaccess.open([result], pqdm_kwargs={"disable": True})
         if len(file_objs) != 1:
@@ -243,9 +265,6 @@ 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", {}))
-
         if xarray_open == "datatree":
             dt = _open_datatree_fn(file_obj, effective_kwargs)
             try:
@@ -283,8 +302,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 +320,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 +352,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 +384,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 +406,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 +420,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 +452,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:

tests/test_plan.py

@@ -2916,7 +2916,7 @@ def test_open_dataset_returns_dataset(
             time_buffer=pd.Timedelta(0),
         )
 
-        ds = p.open_dataset(p[0], open_dataset_kwargs={"engine": "netcdf4"})
+        ds = p.open_dataset(0, open_method={"open_kwargs": {"engine": "netcdf4"}}, silent=True)
         assert isinstance(ds, xr.Dataset)
         assert "sst" in ds
         ds.close()
@@ -2954,7 +2954,7 @@ def test_open_mfdataset_returns_dataset(
         fake_ds = xr.Dataset({"sst": (["lat", "lon"], [[1.0]])}, coords={"lat": [0.0], "lon": [0.0]})
         with patch("xarray.open_mfdataset", return_value=fake_ds) as mock_mfdataset:
             # Pass a list of results directly (backward-compatible path)
-            ds = p.open_mfdataset(fake_results, open_dataset_kwargs={"engine": "netcdf4"})
+            ds = p.open_mfdataset(fake_results, open_method={"open_kwargs": {"engine": "netcdf4"}}, silent=True)
 
         assert ds is fake_ds
         mock_ea.open.assert_called_once_with(fake_results, pqdm_kwargs={"disable": True})
@@ -3016,7 +3016,7 @@ def test_open_mfdataset_accepts_subset_plan(
 
         fake_ds = xr.Dataset({"sst": (["lat", "lon"], [[1.0]])}, coords={"lat": [0.0], "lon": [0.0]})
         with patch("xarray.open_mfdataset", return_value=fake_ds) as mock_mfdataset:
-            ds = p.open_mfdataset(subset, open_dataset_kwargs={"engine": "netcdf4"})
+            ds = p.open_mfdataset(subset, open_method={"open_kwargs": {"engine": "netcdf4"}}, silent=True)
 
         assert ds is fake_ds
         mock_ea.open.assert_called_once_with(fake_results, pqdm_kwargs={"disable": True})
@@ -3047,7 +3047,7 @@ def test_open_dataset_default_uses_auto(
             time_buffer=pd.Timedelta(0),
         )
 
-        ds = p.open_dataset(p[0], open_method="dataset", open_dataset_kwargs={"engine": "netcdf4"})
+        ds = p.open_dataset(0, open_method={"xarray_open": "dataset", "open_kwargs": {"engine": "netcdf4"}}, silent=True)
         assert isinstance(ds, xr.Dataset)
         assert "sst" in ds
         ds.close()
@@ -3073,7 +3073,7 @@ def test_open_dataset_invalid_open_method_raises(
         )
 
         with pytest.raises(ValueError, match="open_method"):
-            p.open_dataset(p[0], open_method="bad")
+            p.open_dataset(0, open_method="bad", silent=True)
 
     def test_open_mfdataset_with_open_method_dataset_uses_open_mfdataset(
         self, tmp_path: pathlib.Path, monkeypatch: pytest.MonkeyPatch
@@ -3106,7 +3106,7 @@ def test_open_mfdataset_with_open_method_dataset_uses_open_mfdataset(
 
         fake_ds = xr.Dataset({"sst": (["lat", "lon"], [[1.0]])}, coords={"lat": [0.0], "lon": [0.0]})
         with patch("xarray.open_mfdataset", return_value=fake_ds) as mock_mfd:
-            ds = p.open_mfdataset(fake_results, open_method="dataset", open_dataset_kwargs={"engine": "netcdf4"})
+            ds = p.open_mfdataset(fake_results, open_method={"xarray_open": "dataset", "open_kwargs": {"engine": "netcdf4"}}, silent=True)
 
         assert ds is fake_ds
         mock_mfd.assert_called_once_with([nc_a, nc_b], chunks={}, engine="netcdf4", decode_timedelta=False)
@@ -3139,14 +3139,123 @@ def test_open_mfdataset_datatree_merge_concatenates(
 
         ds = p.open_mfdataset(
             fake_results,
-            open_method="datatree-merge",
-            open_dataset_kwargs={"engine": "netcdf4"},
+            open_method={"xarray_open": "datatree", "merge": "all", "open_kwargs": {"engine": "netcdf4"}},
+            silent=True,
         )
         # Result is a Dataset with a "granule" dimension from concatenation.
         assert isinstance(ds, xr.Dataset)
         assert ds.sizes["granule"] == 2
         assert "sst" in ds
 
+    def test_open_dataset_prints_effective_spec_when_not_silent(
+        self, tmp_path: pathlib.Path, monkeypatch: pytest.MonkeyPatch, capsys: pytest.CaptureFixture
+    ) -> None:
+        """open_dataset prints the effective open_method spec when silent=False."""
+        nc_path = str(tmp_path / "test.nc")
+        _make_l3_dataset([-90.0, 0.0, 90.0], [-180.0, 0.0, 180.0]).to_netcdf(
+            nc_path, engine="netcdf4"
+        )
+        mock_ea = MagicMock()
+        mock_ea.open.return_value = [nc_path]
+        monkeypatch.setitem(__import__("sys").modules, "earthaccess", mock_ea)
+
+        pts = pd.DataFrame(
+            {"lat": [0.0], "lon": [0.0], "time": pd.to_datetime(["2023-06-01"])}
+        )
+        p = Plan(
+            points=pts,
+            results=[object()],
+            granules=[],
+            point_granule_map={0: []},
+            source_kwargs={"short_name": "TEST"},
+            time_buffer=pd.Timedelta(0),
+        )
+
+        ds = p.open_dataset(0, open_method={"open_kwargs": {"engine": "netcdf4"}}, silent=False)
+        ds.close()
+        captured = capsys.readouterr()
+        assert "open_method:" in captured.out
+        # The effective spec includes the resolved defaults (chunks, decode_timedelta)
+        assert "engine" in captured.out
+
+    def test_open_dataset_silent_suppresses_output(
+        self, tmp_path: pathlib.Path, monkeypatch: pytest.MonkeyPatch, capsys: pytest.CaptureFixture
+    ) -> None:
+        """open_dataset prints nothing when silent=True."""
+        nc_path = str(tmp_path / "test.nc")
+        _make_l3_dataset([-90.0, 0.0, 90.0], [-180.0, 0.0, 180.0]).to_netcdf(
+            nc_path, engine="netcdf4"
+        )
+        mock_ea = MagicMock()
+        mock_ea.open.return_value = [nc_path]
+        monkeypatch.setitem(__import__("sys").modules, "earthaccess", mock_ea)
+
+        pts = pd.DataFrame(
+            {"lat": [0.0], "lon": [0.0], "time": pd.to_datetime(["2023-06-01"])}
+        )
+        p = Plan(
+            points=pts,
+            results=[object()],
+            granules=[],
+            point_granule_map={0: []},
+            source_kwargs={"short_name": "TEST"},
+            time_buffer=pd.Timedelta(0),
+        )
+
+        ds = p.open_dataset(0, open_method={"open_kwargs": {"engine": "netcdf4"}}, silent=True)
+        ds.close()
+        captured = capsys.readouterr()
+        assert captured.out == ""
+
+    def test_open_dataset_integer_index(
+        self, tmp_path: pathlib.Path, monkeypatch: pytest.MonkeyPatch
+    ) -> None:
+        """open_dataset(0) resolves the integer to plan.results[0]."""
+        nc_path = str(tmp_path / "test.nc")
+        _make_l3_dataset([-90.0, 0.0, 90.0], [-180.0, 0.0, 180.0]).to_netcdf(
+            nc_path, engine="netcdf4"
+        )
+        mock_ea = MagicMock()
+        mock_ea.open.return_value = [nc_path]
+        monkeypatch.setitem(__import__("sys").modules, "earthaccess", mock_ea)
+
+        fake_result = object()
+        pts = pd.DataFrame(
+            {"lat": [0.0], "lon": [0.0], "time": pd.to_datetime(["2023-06-01"])}
+        )
+        p = Plan(
+            points=pts,
+            results=[fake_result],
+            granules=[],
+            point_granule_map={0: []},
+            source_kwargs={"short_name": "TEST"},
+            time_buffer=pd.Timedelta(0),
+        )
+
+        # Using integer index 0 should behave identically to plan[0]
+        ds = p.open_dataset(0, open_method={"open_kwargs": {"engine": "netcdf4"}}, silent=True)
+        assert isinstance(ds, xr.Dataset)
+        assert "sst" in ds
+        ds.close()
+        mock_ea.open.assert_called_once_with([fake_result], pqdm_kwargs={"disable": True})
+
+    def test_open_dataset_integer_index_out_of_range_raises(self) -> None:
+        """open_dataset raises IndexError for an out-of-range integer index."""
+        pts = pd.DataFrame(
+            {"lat": [0.0], "lon": [0.0], "time": pd.to_datetime(["2023-06-01"])}
+        )
+        p = Plan(
+            points=pts,
+            results=[object()],
+            granules=[],
+            point_granule_map={0: []},
+            source_kwargs={"short_name": "TEST"},
+            time_buffer=pd.Timedelta(0),
+        )
+
+        with pytest.raises(IndexError, match="out of range"):
+            p.open_dataset(5, silent=True)
+
 
 # ---------------------------------------------------------------------------
 # Helper: create a grouped NetCDF4 file (HDF5 with subgroups)
@@ -3426,7 +3535,7 @@ def test_show_variables_prints_dims_and_vars(
             time_buffer=pd.Timedelta(0),
         )
 
-        p.show_variables(open_dataset_kwargs={"engine": "netcdf4"})
+        p.show_variables(open_method={"open_kwargs": {"engine": "netcdf4"}})
         captured = capsys.readouterr()
         assert "Dimensions" in captured.out
         assert "Variables" in captured.out
@@ -4575,7 +4684,7 @@ def test_show_variables_dataset_layout_prints_dims_and_vars(
             time_buffer=pd.Timedelta(0),
         )
 
-        p.show_variables(open_dataset_kwargs={"engine": "netcdf4"})
+        p.show_variables(open_method={"open_kwargs": {"engine": "netcdf4"}})
         captured = capsys.readouterr()
         assert "Dimensions" in captured.out
         assert "Variables" in captured.out
@@ -4613,7 +4722,7 @@ def test_show_variables_geo_detection_none_warns(
             time_buffer=pd.Timedelta(0),
         )
 
-        p.show_variables(open_dataset_kwargs={"engine": "netcdf4"})
+        p.show_variables(open_method={"open_kwargs": {"engine": "netcdf4"}})
         captured = capsys.readouterr()
         assert "NONE" in captured.out or "no geolocation" in captured.out.lower()