From 22ae309b1a40352b877467135f0c3ba98b704770 Mon Sep 17 00:00:00 2001
From: "Stephen C. Pope" <stephen@descarteslabs.com>
Date: Fri, 9 Jun 2023 11:16:39 -0700
Subject: [PATCH] Shapely multi geometries no longer iterable (#11923)

GitOrigin-RevId: 0143cd182b0b7cc44ccfab0fa505b5c7f31819ff
---
 descarteslabs/core/common/dltile/utm.py       | 17 ++++++++-
 .../core/common/geo/tests/test_geocontext.py  | 35 +++++++++++++++++++
 2 files changed, 51 insertions(+), 1 deletion(-)

diff --git a/descarteslabs/core/common/dltile/utm.py b/descarteslabs/core/common/dltile/utm.py
index e6e45c07..8f43cfd3 100644
--- a/descarteslabs/core/common/dltile/utm.py
+++ b/descarteslabs/core/common/dltile/utm.py
@@ -109,9 +109,24 @@ def _transform(points, *args, axis=-1, **kwargs):
             transformed_points = _transform(points, *args, **kwargs)
             return geo.mapping(transformed_points)
 
+        elif isinstance(points, geo.MultiPoint):
+            return geo.MultiPoint(
+                [_transform(geom, *args, **kwargs) for geom in points.geoms]
+            )
+
+        elif isinstance(points, geo.MultiLineString):
+            return geo.MultiLineString(
+                [_transform(geom, *args, **kwargs) for geom in points.geoms]
+            )
+
         elif isinstance(points, geo.MultiPolygon):
             return geo.MultiPolygon(
-                [_transform(polygon, *args, **kwargs) for polygon in points]
+                [_transform(geom, *args, **kwargs) for geom in points.geoms]
+            )
+
+        elif isinstance(points, geo.GeometryCollection):
+            return geo.GeometryCollection(
+                [_transform(geom, *args, **kwargs) for geom in points.geoms]
             )
 
         elif isinstance(points, Sequence):
diff --git a/descarteslabs/core/common/geo/tests/test_geocontext.py b/descarteslabs/core/common/geo/tests/test_geocontext.py
index 302b9e3b..59c59f00 100644
--- a/descarteslabs/core/common/geo/tests/test_geocontext.py
+++ b/descarteslabs/core/common/geo/tests/test_geocontext.py
@@ -403,6 +403,41 @@ def test_iter_from_shape(self):
         assert type(dltiles1[1]) == str
         assert len(dltiles1) == len(dltiles2)
 
+    def test_iter_from_shape_multi(self):
+        params = {"resolution": 1.5, "tilesize": 512, "pad": 0, "keys_only": True}
+        shape = {
+            "type": "Feature",
+            "geometry": {
+                "type": "MultiPolygon",
+                "coordinates": [
+                    [
+                        [
+                            [-122.51140471760839, 37.77130087547876],
+                            [-122.45475646845254, 37.77475476721895],
+                            [-122.45303985468301, 37.76657207194229],
+                            [-122.51057242081689, 37.763446782666094],
+                            [-122.51140471760839, 37.77130087547876],
+                        ]
+                    ],
+                    [
+                        [
+                            [-123.51140471760839, 37.77130087547876],
+                            [-123.45475646845254, 37.77475476721895],
+                            [-123.45303985468301, 37.76657207194229],
+                            [-123.51057242081689, 37.763446782666094],
+                            [-123.51140471760839, 37.77130087547876],
+                        ]
+                    ],
+                ],
+            },
+            "properties": None,
+        }
+        dltiles1 = [tile for tile in geocontext.DLTile.iter_from_shape(shape, **params)]
+        dltiles2 = geocontext.DLTile.from_shape(shape, **params)
+        assert type(dltiles1[0]) == str
+        assert type(dltiles1[1]) == str
+        assert len(dltiles1) == len(dltiles2)
+
     def test_subtile(self):
         tile = geocontext.DLTile.from_key("2048:0:30.0:15:3:80")
         tiles = [t for t in tile.subtile(8, keys_only=True)]