diff --git a/xrspatial/polygonize.py b/xrspatial/polygonize.py index 6bf4a3ac..7f5d76d1 100644 --- a/xrspatial/polygonize.py +++ b/xrspatial/polygonize.py @@ -45,7 +45,12 @@ except ImportError: cupy = None -from .utils import ArrayTypeFunctionMapping, is_cupy_array, ngjit +from .utils import ( + ArrayTypeFunctionMapping, + _validate_raster, + is_cupy_array, + ngjit, +) _regions_dtype = np.uint32 _visited_dtype = np.uint8 @@ -1620,12 +1625,15 @@ def polygonize( For Dask+CuPy, each chunk is transferred independently, keeping peak CPU memory proportional to chunk size rather than full raster size. """ - if raster.ndim != 2 or raster.shape[0] < 1 or raster.shape[1] < 1: + _validate_raster(raster, func_name='polygonize', name='raster', ndim=2) + if raster.shape[0] < 1 or raster.shape[1] < 1: raise ValueError( "Raster array must be 2D with a shape of at least (1, 1)") # Check mask. if mask is not None: + _validate_raster(mask, func_name='polygonize', name='mask', + ndim=2, numeric=False) if not (type(raster.data) is type(mask.data)): # noqa: E721 raise TypeError( "raster and mask have different underlying types: " diff --git a/xrspatial/tests/test_polygonize.py b/xrspatial/tests/test_polygonize.py index 0ad2c01f..07ead34f 100644 --- a/xrspatial/tests/test_polygonize.py +++ b/xrspatial/tests/test_polygonize.py @@ -1068,3 +1068,49 @@ def test_simplify_cupy_matches_numpy(self): for val in areas_np: assert_allclose(areas_cp[val], areas_np[val], atol=1e-10) + + +# ===================================================================== +# Issue #1441: _validate_raster on raster + mask +# ===================================================================== + +import xarray as _xr_for_validation + + +class TestPolygonizeInputValidation: + """polygonize() rejects bad raster / mask inputs (#1441).""" + + @staticmethod + def _good_raster(): + return _xr_for_validation.DataArray( + np.zeros((4, 4), dtype=np.int32), + dims=('y', 'x'), + coords={'y': np.arange(4), 'x': np.arange(4)}, + ) + + def test_rejects_non_dataarray_raster(self): + from xrspatial.polygonize import polygonize + with pytest.raises(TypeError, match="xarray.DataArray"): + polygonize(np.zeros((4, 4), dtype=np.int32)) + + def test_rejects_complex_dtype_raster(self): + from xrspatial.polygonize import polygonize + bad = _xr_for_validation.DataArray( + np.zeros((4, 4), dtype=np.complex128), + dims=('y', 'x'), + coords={'y': np.arange(4), 'x': np.arange(4)}, + ) + with pytest.raises(ValueError, match="real numeric"): + polygonize(bad) + + def test_rejects_1d_raster(self): + from xrspatial.polygonize import polygonize + bad = _xr_for_validation.DataArray( + np.zeros(4, dtype=np.int32), dims=('y',)) + with pytest.raises(ValueError, match=r"must be 2D"): + polygonize(bad) + + def test_rejects_non_dataarray_mask(self): + from xrspatial.polygonize import polygonize + with pytest.raises(TypeError, match="xarray.DataArray"): + polygonize(self._good_raster(), mask=np.ones((4, 4), dtype=bool))