Merge pull request #290 from ericboucher/augment-custom-functions

Add rv_array to custom functions

Author: perrygeo

src/rasterstats/main.py

@@ -1,3 +1,4 @@
+import inspect
 import sys
 import warnings
 
@@ -283,10 +284,14 @@ def gen_zonal_stats(
 
             if add_stats is not None:
                 for stat_name, stat_func in add_stats.items():
-                    try:
+                    n_params = len(inspect.signature(stat_func).parameters.keys())
+                    if n_params == 3:
+                        feature_stats[stat_name] = stat_func(masked, feat["properties"], rv_array)
+                    # backwards compatible with two-argument function
+                    elif n_params == 2:
                         feature_stats[stat_name] = stat_func(masked, feat["properties"])
-                    except TypeError:
-                        # backwards compatible with single-argument function
+                    # backwards compatible with single-argument function
+                    else:
                         feature_stats[stat_name] = stat_func(masked)
 
             if raster_out:

tests/test_zonal.py

@@ -315,6 +315,18 @@ def mymean_prop(x, prop):
     for i in range(len(stats)):
         assert stats[i]["mymean_prop"] == stats[i]["mean"] * (i + 1)
 
+def test_add_stats_prop_and_array():
+    polygons = os.path.join(DATA, "polygons.shp")
+
+    def mymean_prop_and_array(x, prop, rv_array):
+        # confirm that the object exists and is accessible.
+        assert rv_array is not None
+        return np.ma.mean(x) * prop["id"]
+
+    stats = zonal_stats(polygons, raster, add_stats={"mymean_prop_and_array": mymean_prop_and_array})
+    for i in range(len(stats)):
+        assert stats[i]["mymean_prop_and_array"] == stats[i]["mean"] * (i + 1)
+
 
 def test_mini_raster():
     polygons = os.path.join(DATA, "polygons.shp")