Merge pull request #203 from drnextgis/zone_func

Add support return statement to zone_func

src/rasterstats/main.py

@@ -188,7 +188,11 @@ def gen_zonal_stats(
                     raise TypeError(('zone_func must be a callable '
                                      'which accepts function a '
                                      'single `zone_array` arg.'))
-                zone_func(masked)
+                value = zone_func(masked)
+
+                # check if zone_func has return statement
+                if value is not None:
+                    masked = value
 
             if masked.compressed().size == 0:
                 # nothing here, fill with None and move on

tests/test_io.py

@@ -147,6 +147,13 @@ def test_jsonstr_collection():
     _test_read_features(indata)
 
 
+def test_jsonstr_collection_without_features():
+    indata = {'type': "FeatureCollection", 'features': []}
+    indata = json.dumps(indata)
+    with pytest.raises(ValueError):
+        _test_read_features(indata)
+
+
 def test_invalid_jsonstr():
     indata = {'type': "InvalidGeometry", 'coordinates': [30, 10]}
     indata = json.dumps(indata)

tests/test_zonal.py

@@ -327,6 +327,20 @@ def test_percentile_good():
     assert stats[0]['percentile_50'] <= stats[0]['percentile_90']
 
 
+def test_zone_func_has_return():
+
+    def example_zone_func(zone_arr):
+        return np.ma.masked_array(np.full(zone_arr.shape, 1))
+
+    polygons = os.path.join(DATA, 'polygons.shp')
+    stats = zonal_stats(polygons,
+                        raster,
+                        zone_func=example_zone_func)
+    assert stats[0]['max'] == 1
+    assert stats[0]['min'] == 1
+    assert stats[0]['mean'] == 1
+
+
 def test_zone_func_good():
 
     def example_zone_func(zone_arr):