Merge e2c93df into ee49927

perrygeo · Dec 1, 2019 · e32953c · e32953c
2 parents ee49927 + e2c93df
commit e32953c
Show file tree

Hide file tree

Showing 2 changed files with 19 additions and 1 deletion.
diff --git a/src/rasterstats/main.py b/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

diff --git a/tests/test_zonal.py b/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):