Merge pull request #202 from jbouffard/mask-fix-2

Mask Can Now Take Either a List or Just a Single Item
locationtech-labs · May 15, 2017 · 7f0b1de · 7f0b1de
2 parents 9732ee2 + 9572a8d
commit 7f0b1de
Show file tree

Hide file tree

Showing 2 changed files with 4 additions and 1 deletion.
diff --git a/geopyspark/geotrellis/rdd.py b/geopyspark/geotrellis/rdd.py
@@ -665,6 +665,9 @@ def mask(self, geometries):
         Returns:
             :class:`~geopyspark.geotrellis.rdd.TiledRasterRDD`
         """
+
+        if not isinstance(geometries, list):
+            geometries = [geometries]
         wkts = [shapely.wkt.dumps(g) for g in geometries]
         srdd = self.srdd.mask(wkts)
 

diff --git a/geopyspark/tests/mask_test.py b/geopyspark/tests/mask_test.py
@@ -38,7 +38,7 @@ class MaskTest(BaseTestClass):
                     'extent': extent,
                     'tileLayout': layout}}
 
-    geometries = [Polygon([(17, 17), (42, 17), (42, 42), (17, 42)])]
+    geometries = Polygon([(17, 17), (42, 17), (42, 42), (17, 42)])
     raster_rdd = TiledRasterRDD.from_numpy_rdd(BaseTestClass.geopysc, SPATIAL, rdd, metadata)
 
     @pytest.fixture(autouse=True)