fix: address review comments from PR #124 (utils cleanup, accessor behavior, tests) (#130)

* test: cover utils helpers and xsg accessor; fix deprecations

- Add tests for format_bytes, asdatetime, compute_2d_subset_mask, and
  utils.assign_ugrid_topology forwarding with DeprecationWarning
- Add accessor tests for unknown-grid warning and None subset behavior
- Fix warnings.warn argument order in assign_ugrid_topology; correct message
- Coerce polygon coords via np.asarray in normalize_polygon_x_coords
- Define zarr__version__ when fsspec/zarr import fails (test collection)

Made-with: Cursor

* fix: address PR #124 review (subset_vars, utils cleanup, tests)

- Remove deprecated assign_ugrid_topology from utils (use grids.ugrid only)

- Raise ValueError from xsg.subset_vars when no grid is recognized

- test_utils: single utils import style; stronger mask checks; list polygon test

- test_accessor: expect ValueError for subset_vars without grid

- test_sgrid: drop zarr>=3 skip for online kerchunk test (per review)

Made-with: Cursor

Author: Divyateja2709

tests/test_accessor.py

@@ -0,0 +1,40 @@
+import numpy as np
+import pytest
+import xarray as xr
+
+import xarray_subset_grid.accessor  # noqa: F401 -- register accessor
+
+
+def test_accessor_warns_when_no_grid_recognized():
+    ds = xr.Dataset()
+    with pytest.warns(UserWarning, match="no grid type"):
+        accessor = ds.xsg
+    assert accessor.grid is None
+
+
+def test_subset_polygon_and_bbox_return_none_without_grid():
+    ds = xr.Dataset()
+    poly = np.array(
+        [
+            [-72.0, 41.0],
+            [-70.0, 41.0],
+            [-71.0, 39.0],
+            [-72.0, 41.0],
+        ]
+    )
+    with pytest.warns(UserWarning, match="no grid type"):
+        assert ds.xsg.subset_polygon(poly) is None
+        assert ds.xsg.subset_bbox((-72, 39, -70, 41)) is None
+
+
+def test_subset_vars_raises_without_grid():
+    ds = xr.Dataset({"a": (("x",), [1, 2, 3])})
+    with pytest.warns(UserWarning, match="no grid type"):
+        with pytest.raises(ValueError, match="subset_vars requires a recognized grid"):
+            ds.xsg.subset_vars(["a"])
+
+
+def test_has_vertical_levels_false_without_grid():
+    ds = xr.Dataset()
+    with pytest.warns(UserWarning, match="no grid type"):
+        assert ds.xsg.has_vertical_levels is False

tests/test_grids/test_sgrid.py

@@ -10,9 +10,6 @@
 # open dataset as zarr object using fsspec reference file system and xarray
 try:
     import fsspec
-    import zarr
-
-    zarr__version__ = int(zarr.__version__.split(".")[0])
 except ImportError:
     fsspec = None
 
@@ -51,9 +48,6 @@ def test_grid_topology_location_parse():
     }
 
 
-@pytest.mark.skipif(
-    zarr__version__ >= 3, reason="zarr3.0.8 doesn't support FSpec AWS (it might soon)"
-)
 @pytest.mark.online
 def test_polygon_subset():
     """

tests/test_utils.py

@@ -1,13 +1,12 @@
 import os
+from datetime import datetime
 
+import cftime
 import numpy as np
 import pytest
+import xarray as xr
 
-from xarray_subset_grid.utils import (
-    normalize_bbox_x_coords,
-    normalize_polygon_x_coords,
-    ray_tracing_numpy,
-)
+import xarray_subset_grid.utils as xsg_utils
 
 # normalize_polygon_x_coords tests.
 
@@ -69,12 +68,7 @@ def get_test_file_dir():
 )
 def test_normalize_x_coords(lons, poly, norm_poly):
     lons = np.array(lons)
-    normalized_polygon = normalize_polygon_x_coords(lons, np.array(poly))
-    print(f"{lons=}")
-    print(f"{poly=}")
-    print(f"{norm_poly=}")
-    print(f"{normalized_polygon=}")
-
+    normalized_polygon = xsg_utils.normalize_polygon_x_coords(lons, np.array(poly))
     assert np.allclose(normalized_polygon, norm_poly)
 
 
@@ -97,7 +91,7 @@ def test_normalize_x_coords(lons, poly, norm_poly):
 )
 def test_normalize_x_coords_bbox(lons, bbox, norm_bbox):
     lons = np.array(lons)
-    normalized_polygon = normalize_bbox_x_coords(lons, bbox)
+    normalized_polygon = xsg_utils.normalize_bbox_x_coords(lons, bbox)
     assert np.allclose(normalized_polygon, norm_bbox)
 
 
@@ -122,6 +116,82 @@ def test_ray_tracing_numpy():
         ]
     )
 
-    result = ray_tracing_numpy(points[:, 0], points[:, 1], poly)
+    result = xsg_utils.ray_tracing_numpy(points[:, 0], points[:, 1], poly)
 
     assert np.array_equal(result, [False, True, False])
+
+
+@pytest.mark.parametrize(
+    "num, unit",
+    [
+        (512, "bytes"),
+        (2048, "KB"),
+        (3 * 1024**2, "MB"),
+    ],
+)
+def test_format_bytes(num, unit):
+    assert unit in xsg_utils.format_bytes(num)
+
+
+def test_asdatetime_none():
+    assert xsg_utils.asdatetime(None) is None
+
+
+def test_asdatetime_datetime_passthrough():
+    dt = datetime(2020, 6, 15, 12, 30, 0)
+    assert xsg_utils.asdatetime(dt) is dt
+
+
+def test_asdatetime_cftime_passthrough():
+    dt = cftime.datetime(2020, 6, 15, 12)
+    assert xsg_utils.asdatetime(dt) is dt
+
+
+def test_asdatetime_parse_string():
+    dt = xsg_utils.asdatetime("2020-06-15T12:30:00")
+    assert dt.year == 2020 and dt.month == 6 and dt.day == 15
+
+
+def test_compute_2d_subset_mask_all_inside():
+    ny, nx = 5, 5
+    lat = np.linspace(40.0, 44.0, ny)
+    lon = np.linspace(-74.0, -70.0, nx)
+    lat2d, lon2d = np.meshgrid(lat, lon, indexing="ij")
+    lat_da = xr.DataArray(lat2d, dims=("y", "x"))
+    lon_da = xr.DataArray(lon2d, dims=("y", "x"))
+    poly = np.array([(-75.0, 39.0), (-69.0, 39.0), (-69.0, 45.0), (-75.0, 45.0)])
+    mask = xsg_utils.compute_2d_subset_mask(lat_da, lon_da, poly)
+    assert mask.dims == ("y", "x")
+    assert mask.all()
+
+
+def test_compute_2d_subset_mask_partial():
+    # Include explicit lon/lat nodes inside the polygon so the mask can be checked at a
+    # non-boundary grid point (ray-casting is ambiguous on polygon edges).
+    lat = np.array([40.0, 40.5, 41.0, 43.0, 46.0])
+    lon = np.array([-74.5, -73.75, -73.0, -71.0, -68.0])
+    lat2d, lon2d = np.meshgrid(lat, lon, indexing="ij")
+    lat_da = xr.DataArray(lat2d, dims=("y", "x"))
+    lon_da = xr.DataArray(lon2d, dims=("y", "x"))
+    poly = np.array([(-74.5, 40.0), (-73.0, 40.0), (-73.0, 41.0), (-74.5, 41.0)])
+    mask = xsg_utils.compute_2d_subset_mask(lat_da, lon_da, poly)
+    assert mask.dims == ("y", "x")
+    assert mask.any()
+    assert not mask.all()
+    i_inside = int(np.where(lat == 40.5)[0][0])
+    j_inside = int(np.where(lon == -73.75)[0][0])
+    assert mask.values[i_inside, j_inside]
+    assert not mask.values[-1, -1]
+
+
+def test_compute_2d_subset_mask_list_polygon_coerced():
+    """list/tuple polygon vertices are accepted (coerced via normalize_polygon_x_coords)."""
+    ny, nx = 5, 5
+    lat = np.linspace(40.0, 44.0, ny)
+    lon = np.linspace(-74.0, -70.0, nx)
+    lat2d, lon2d = np.meshgrid(lat, lon, indexing="ij")
+    lat_da = xr.DataArray(lat2d, dims=("y", "x"))
+    lon_da = xr.DataArray(lon2d, dims=("y", "x"))
+    poly = [(-75.0, 39.0), (-69.0, 39.0), (-69.0, 45.0), (-75.0, 45.0)]
+    mask = xsg_utils.compute_2d_subset_mask(lat_da, lon_da, poly)
+    assert mask.all()

xarray_subset_grid/accessor.py

@@ -112,9 +112,12 @@ def subset_vars(self, vars: list[str]) -> xr.Dataset:
         :param vars: The variables to keep
         :return: The subsetted dataset
         """
-        if self._grid:
-            return self._grid.subset_vars(self._ds, vars)
-        return self._ds
+        if not self._grid:
+            raise ValueError(
+                "subset_vars requires a recognized grid; this dataset has no grid type "
+                "that xarray-subset-grid can use."
+            )
+        return self._grid.subset_vars(self._ds, vars)
 
     @property
     def has_vertical_levels(self) -> bool:

xarray_subset_grid/utils.py

@@ -1,4 +1,3 @@
-import warnings
 from datetime import datetime
 
 import cf_xarray  # noqa
@@ -19,12 +18,14 @@ def normalize_polygon_x_coords(x, poly):
     If the x coords are between 0 and 180 (i.e. both will work), the polygon
     is not changed.
 
-    NOTE: polygon is normalized in place!
+    NOTE: ``poly`` is normalized in place when it is already an ndarray;
+    a copy is made when ``poly`` is a sequence.
 
     Args:
         x (np.array): x-coordinates of the vertices
         poly (np.array): polygon vertices
     """
+    poly = np.asarray(poly)
     x_min, x_max = x.min(), x.max()
 
     poly_x = poly[:, 0]
@@ -93,20 +94,6 @@ def ray_tracing_numpy(x, y, poly):
     return inside
 
 
-# This is defined in ugrid.py
-# this placeholder for backwards compatibility for a brief period
-def assign_ugrid_topology(*args, **kwargs):
-    warnings.warn(
-        DeprecationWarning,
-        "The function `assign_grid_topology` has been moved to the "
-        "`grids.ugrid` module. It will not be able to be called from "
-        "the utils `module` in the future.",
-    )
-    from .grids.ugrid import assign_ugrid_topology
-
-    return assign_ugrid_topology(*args, **kwargs)
-
-
 def format_bytes(num):
     """This function will convert bytes to MB....