Add new crop method for cropping xr data to geometries

odc/geo/_xr_interop.py

@@ -245,6 +245,59 @@ def assign_crs(
     return xx
 
 
+def crop(xx: XrT, poly: Geometry, mask: bool = True, all_touched: bool = True) -> XrT:
+    """
+    Crops and optionally mask an xr.Dataset or xr.DataArray (``xx``) to 
+    the spatial extent of a geometry (``poly``).
+
+    :param xx:
+       :py:class:`~xarray.Dataset` or :py:class:`~xarray.DataArray`.
+
+    :param poly:
+       Geometry shape used to crop ``xx``.
+
+    :param poly:
+       Whether to mask out pixels outside of the rasterized extent of 
+       ``poly``.
+
+    :param all_touched:
+        If ``True``, all pixels touched by ``poly`` will remain unmasked.
+        If ``False``, only pixels whose center is within the polygon or 
+        that are selected by Bresenham's line algorithm will remain unmasked.
+
+    :return:
+        A :py:class:`~xarray.Dataset` or :py:class:`~xarray.DataArray`
+        cropped and optionally masked to the spatial extent of ``poly``.
+    """
+    # Reproject `poly` to match `xx` and calculate bounding box
+    poly = poly.to_crs(crs=xx.odc.crs)
+    bbox = poly.boundingbox
+
+    # Verify that `poly` overlaps with `xx` extent
+    if poly.intersects(xx.odc.geobox.extent):
+        # Crop `xx` to the bounding box of `poly`
+        ydim, xdim = spatial_dims(xx)
+        xx_cropped = xx.sel(
+            {xdim: slice(bbox.left, bbox.right), ydim: slice(bbox.top, bbox.bottom)}
+        )
+
+        # Optionally cropped data outside rasterized `poly`
+        if mask:
+            rasterized = rasterize(
+                poly=poly,
+                how=xx_cropped.odc.geobox,
+                all_touched=all_touched,
+            )
+            xx_cropped = xx_cropped.where(rasterized)
+
+        return xx_cropped
+
+    else:
+        raise Exception(
+            "The supplied `poly` must overlap spatially with the extent of `xx`."
+        )
+
+
 def xr_coords(
     gbox: SomeGeoBox, crs_coord_name: Optional[str] = _DEFAULT_CRS_COORD_NAME
 ) -> Dict[Hashable, xarray.DataArray]:
@@ -731,6 +784,7 @@ def nodata(self) -> Optional[float]:
         return None
 
     colorize = _wrap_op(colorize)
+    crop = _wrap_op(crop)
 
     if have.rasterio:
         write_cog = _wrap_op(write_cog)
@@ -767,6 +821,8 @@ def to_rgba(
         vmax: Optional[float] = None,
     ) -> xarray.DataArray:
         return to_rgba(self._xx, bands=bands, vmin=vmin, vmax=vmax)
+
+    crop = _wrap_op(crop)
 
     if have.rasterio:
         reproject = _wrap_op(_xr_reproject_ds)

odc/geo/xr.py

@@ -18,6 +18,7 @@
     ODCExtensionDs,
     assign_crs,
     colorize,
+    crop,
     rasterize,
     register_geobox,
     spatial_dims,
@@ -45,6 +46,7 @@
     "xr_zeros",
     "colorize",
     "to_rgba",
+    "crop",
 ]
 
 # pylint: disable=import-outside-toplevel,unused-import