feat(helpers): accept & auto-convert ee.Geometry in fit_geometry

docs/faq.md

@@ -29,7 +29,7 @@ To align with CF conventions (`[time, y, x]`) and reduce the need for transposes
 Your AOI may fall outside the dataset extent or the CRS mismatch caused an unexpected reprojection. Try matching source grid first to confirm availability.
 
 ## Do I need shapely geometries?
-Helpers accept shapely for convenience. If you already have an EE geometry, you can convert it to shapely with `shapely.geometry.shape(ee_geom.getInfo())`. Shapely makes reprojection and area reasoning simpler client-side.
+Helpers accept shapely for convenience, as it makes reprojection and area reasoning simpler client-side. However, `fit_geometry` also accepts an Earth Engine geometry (e.g. `ee.Geometry`) directly and converts it for you via `shapely.geometry.shape(ee_geom.getInfo())`. That said, if you do need a shapely geometry elsewhere, you may want to handle this conversion explicitly.
 
 ## `ds.to_netcdf()` fails with `ValueError: could not safely cast array from int64 to int32`
 Xee time coordinates are stored as `int64` (nanoseconds since epoch). The `scipy` netCDF writer only supports netCDF3, which is limited to `int32`, so the write fails when `scipy` is the only available backend.

docs/guide.md

@@ -57,6 +57,40 @@ grid_params = helpers.fit_geometry(
 ds = xr.open_dataset('ee://ECMWF/ERA5_LAND/MONTHLY_AGGR', engine='ee', **grid_params)
 ```
 
+## Using an Earth Engine Geometry as the AOI
+
+`fit_geometry` accepts an `ee.Geometry` directly. It is auto-converted to shapely via `shapely.geometry.shape(geometry.getInfo())`, so you can stay in an Earth Engine-native workflow without converting by hand.
+
+```python
+import ee
+import xarray as xr
+from xee import helpers
+
+aoi_ee = ee.Geometry.Rectangle([113.33, -43.63, 153.56, -10.66])
+grid_params = helpers.fit_geometry(
+    geometry=aoi_ee,
+    grid_crs='EPSG:4326',
+    grid_shape=(256, 256),
+)
+
+ds = xr.open_dataset('ee://ECMWF/ERA5_LAND/MONTHLY_AGGR', engine='ee', **grid_params)
+```
+
+If you prefer to convert explicitly or need the shapely geometry elsewhere, the equivalent end-to-end conversion is as follows:
+
+```python
+import shapely
+
+aoi_shapely = shapely.geometry.shape(aoi_ee.getInfo())
+grid_params = helpers.fit_geometry(
+    geometry=aoi_shapely,
+    grid_crs='EPSG:4326',
+    grid_shape=(256, 256),
+)
+```
+
+Non-geometry inputs (or geometries that can't be converted) raise a `TypeError` that includes the above conversion snippet.
+
 ## Custom Region at Source Resolution
 
 Fit an AOI but keep original pixel size.

docs/quickstart.md

@@ -123,6 +123,12 @@ grid = helpers.fit_geometry(
 ds = xr.open_dataset('ee://ECMWF/ERA5_LAND/MONTHLY_AGGR', engine='ee', **grid)
 ```
 
+```{admonition} Already have an Earth Engine geometry?
+:class: tip
+
+`helpers.fit_geometry` accepts an `ee.Geometry` directly. See the [User Guide](guide.md) for an end-to-end example.
+```
+
 ## 7. Having trouble?
 
 See the [FAQ](faq.md) and open a [discussion](https://github.com/google/Xee/discussions) if needed.

xee/ext_test.py

@@ -384,6 +384,64 @@ def test_fit_geometry_with_rounding(self):
     self.assertAlmostEqual(grid_dict['crs_transform'][0], 0.1)
     self.assertAlmostEqual(grid_dict['crs_transform'][4], -0.1)
 
+  def test_fit_geometry_accepts_ee_like_geometry(self):
+    """Test that an Earth Engine-like geometry is auto-converted via getInfo()."""
+
+    class _FakeEEGeometry:
+      """Minimal ee.Geometry stand-in exposing getInfo()."""
+
+      def __init__(self, geojson):
+        self._geojson = geojson
+
+      def getInfo(self):
+        return self._geojson
+
+    ee_like_geometry = _FakeEEGeometry(
+        {
+            'type': 'Polygon',
+            'coordinates': [[[10.1, 10.1], [10.1, 10.9], [11.9, 10.1]]],
+        }
+    )
+
+    grid_dict = helpers.fit_geometry(
+        geometry=ee_like_geometry,
+        grid_crs='EPSG:4326',
+        grid_scale=(0.5, -0.5),
+    )
+
+    self.assertEqual(
+        grid_dict['crs_transform'], (0.5, 0.0, 10.0, 0.0, -0.5, 11.0)
+    )
+    self.assertEqual(grid_dict['shape_2d'], (4, 2))
+
+  def test_fit_geometry_ee_like_invalid_payload_raises(self):
+    """Test that an Earth Engine-like geometry with an invalid payload raises TypeError."""
+
+    class _FakeEEGeometry:
+
+      def getInfo(self):
+        return {'type': 'not a geojson geometry'}
+
+    with self.assertRaisesRegex(
+        TypeError, r'shapely\.geometry\.shape\(ee_geom\.getInfo\(\)\)'
+    ):
+      helpers.fit_geometry(
+          geometry=_FakeEEGeometry(),
+          grid_crs='EPSG:4326',
+          grid_scale=(0.5, -0.5),
+      )
+
+  def test_fit_geometry_unsupported_type_raises(self):
+    """Test that an unrelated input type raises a clear TypeError with guidance."""
+    with self.assertRaisesRegex(
+        TypeError, r'shapely\.geometry\.shape\(ee_geom\.getInfo\(\)\)'
+    ):
+      helpers.fit_geometry(
+          geometry=42,
+          grid_crs='EPSG:4326',
+          grid_scale=(0.5, -0.5),
+      )
+
 
 if __name__ == '__main__':
   absltest.main()

xee/helpers.py

@@ -98,9 +98,61 @@ def set_scale(
   affine_transform = affine.Affine(*crs_transform)
   return list(affine_transform)[:6]
 
+def _coerce_to_shapely_geometry(
+    geometry: Union[shapely.geometry.base.BaseGeometry, ee.Geometry],
+) -> shapely.geometry.base.BaseGeometry:
+    """Normalize a supported geometry input to a shapely geometry.
+
+    Shapely geometries are returned unchanged. Earth Engine-like geometries are
+    automatically detected and converted. Any other input raises a ``TypeError``
+    that names the expected type and includes the explicit conversion snippet.
+
+    Args:
+      geometry: A shapely geometry or an Earth Engine-like geometry exposing
+        ``getInfo``.
+
+    Returns:
+      An equivalent shapely geometry.
+
+    Raises:
+      TypeError: If ``geometry`` is neither a shapely geometry nor convertible
+        from an Earth Engine-like geometry.
+    """
+    if isinstance(geometry, shapely.geometry.base.BaseGeometry):
+      return geometry
+
+    get_info = getattr(geometry, "getInfo", None)
+
+    if callable(get_info):
+      # NOTE(abi): ``getInfo`` runs outside the try clock so that genuine EE
+      #            runtime errors propagate unchanged.
+      geojson = get_info()
+
+      try:
+        return shapely.geometry.shape(geojson)
+      except (
+          AttributeError,
+          KeyError,
+          TypeError,
+          ValueError,
+          shapely.errors.GeometryTypeError,
+      ) as e:
+          raise TypeError(
+              "Could not convert the Earth Engine-like geometry to a shapely "
+              "geometry. Convert it explicitly before calling fit_geometry:\n"
+              "    shapely.geometry.shape(ee_geom.getInfo())"
+          ) from e
+
+    raise TypeError(
+        "fit_geometry expected a shapely geometry, but got "
+        f"{type(geometry).__name__!r}. If this is an Earth Engine geometry, "
+        "convert it with:\n"
+        "    shapely.geometry.shape(ee_geom.getInfo())"
+    )
+
 
 def fit_geometry(
-    geometry: shapely.geometry.base.BaseGeometry,
+    geometry: Union[shapely.geometry.base.BaseGeometry, ee.Geometry],
     # All following parameters are keyword-only.
     *,
     geometry_crs: str = 'EPSG:4326',
@@ -119,7 +171,8 @@ def fit_geometry(
 
   Args:
     geometry: Shapely geometry defining the area of interest (in
-      ``geometry_crs`` units).
+      ``geometry_crs`` units). An Earth Engine-like geometry exposing
+      ``getInfo`` is also accepted and converted automatically.
     geometry_crs: CRS of the input geometry (default WGS84).
     buffer: Optional positive distance in CRS units to expand the geometry.
     grid_crs: Target CRS for the output grid.
@@ -142,6 +195,8 @@ def fit_geometry(
         "Exactly one of 'grid_scale' or 'grid_shape' must be specified."
     )
 
+  geometry = _coerce_to_shapely_geometry(geometry)
+
   transformer = Transformer.from_crs(
       crs_from=geometry_crs, crs_to=grid_crs, always_xy=True
   )