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

docs/faq.md

@@ -29,7 +29,16 @@ 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.
+`fit_geometry` accepts **shapely geometries** or **ee.Geometry** only. Shapely is convenient for client-side reprojection and area reasoning. If you have an `ee.Geometry`, it is auto-converted for you via `shapely.geometry.shape(ee_geom.getInfo())`.
+
+If you have an `ee.Feature` or other `ee.ComputedObject`, call `.geometry()` first:
+```python
+# Correct
+fit_geometry(geometry=ee_feature.geometry(), ...)
+
+# Incorrect
+fit_geometry(geometry=ee_feature, ...)
+```
 
 ## `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,55 @@ 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 `ee.Geometry` directly and auto-converts it to shapely via `shapely.geometry.shape(geometry.getInfo())`.
+
+```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)
+```
+
+**Important:** If you have an `ee.Feature` or other `ee.ComputedObject`, extract the geometry first:
+
+```python
+aoi_feature = ee.Feature(...)
+# Correct - Extracting the geometry
+grid_params = helpers.fit_geometry(
+    geometry=aoi_feature.geometry(),
+    grid_crs='EPSG:4326',
+    grid_shape=(256, 256),
+)
+
+# Incorrect - Passing the feature directly
+grid_params = helpers.fit_geometry(geometry=aoi_feature, ...)
+```
+
+If you prefer to convert to shapely explicitly or need the shapely geometry elsewhere:
+
+```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 `shapely.geometry` or `ee.Geometry` directly. If you have an `ee.Feature`, call `.geometry()` first. See the [User Guide](guide.md) for examples.
+```
+
 ## 7. Having trouble?
 
 See the [FAQ](faq.md) and open a [discussion](https://github.com/google/Xee/discussions) if needed.

xee/ext_test.py

@@ -4,6 +4,7 @@
 from absl.testing import parameterized
 import numpy as np
 import affine
+import ee
 import shapely
 from unittest import mock
 import xee
@@ -384,6 +385,50 @@ 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_geometry(self):
+    """Test that an ee.Geometry is auto-converted via getInfo()."""
+
+    ee_geometry = mock.MagicMock(spec=ee.Geometry)
+    ee_geometry.getInfo.return_value = {
+        'type': 'Polygon',
+        'coordinates': [[[10.1, 10.1], [10.1, 10.9], [11.9, 10.1]]],
+    }
+
+    grid_dict = helpers.fit_geometry(
+        geometry=ee_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_geometry_invalid_payload_raises(self):
+    """Test that an ee.Geometry with an invalid payload raises TypeError."""
+
+    ee_geometry = mock.MagicMock(spec=ee.Geometry)
+    ee_geometry.getInfo.return_value = {'type': 'not a geojson geometry'}
+
+    with self.assertRaisesRegex(
+        TypeError, r'shapely\.geometry\.shape\(ee_geom\.getInfo\(\)\)'
+    ):
+      helpers.fit_geometry(
+          geometry=ee_geometry,
+          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 or ee\.Geometry'):
+      helpers.fit_geometry(
+          geometry=42,
+          grid_crs='EPSG:4326',
+          grid_scale=(0.5, -0.5),
+      )
+
 
 if __name__ == '__main__':
   absltest.main()

xee/helpers.py

@@ -99,8 +99,57 @@ def set_scale(
   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. ee.Geometry inputs are
+  automatically 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 ee.Geometry instance.
+
+  Returns:
+    An equivalent shapely geometry.
+
+  Raises:
+    TypeError: If ``geometry`` is neither a shapely geometry nor an ee.Geometry.
+  """
+  if isinstance(geometry, shapely.geometry.base.BaseGeometry):
+    return geometry
+
+  if isinstance(geometry, ee.Geometry):
+    # NOTE: ``getInfo`` runs outside the try block so that genuine EE
+    #       runtime errors propagate unchanged.
+    geojson = geometry.getInfo()
+
+    try:
+      return shapely.geometry.shape(geojson)
+    except (
+        AttributeError,
+        KeyError,
+        TypeError,
+        ValueError,
+        shapely.errors.GeometryTypeError,
+    ) as e:
+      raise TypeError(
+          "Could not convert the ee.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 or ee.Geometry, but got "
+      f"{type(geometry).__name__!r}. If using an ee.Feature or other "
+      "ee.ComputedObject, convert it explicitly 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',
@@ -118,8 +167,8 @@ def fit_geometry(
   provided the scale is inferred uniformly over the geometry's bounding box.
 
   Args:
-    geometry: Shapely geometry defining the area of interest (in
-      ``geometry_crs`` units).
+    geometry: Shapely geometry or ee.Geometry defining the area of interest (in
+      ``geometry_crs`` units). An ee.Geometry is 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 +191,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
   )