From e2062459741a84a9e3f4ac17387aa8c2fbb80ae6 Mon Sep 17 00:00:00 2001
From: msoltadeo <msoltadeo@gmail.com>
Date: Mon, 25 Nov 2019 16:42:10 -0300
Subject: [PATCH] add the how_fillna option

---
 variable_generators/generators.py | 12 ++++++++----
 1 file changed, 8 insertions(+), 4 deletions(-)

diff --git a/variable_generators/generators.py b/variable_generators/generators.py
index 5157c3a..0592e09 100644
--- a/variable_generators/generators.py
+++ b/variable_generators/generators.py
@@ -14,7 +14,7 @@
     pass
 
 
-def make_agg_var(agent, geog, geog_id, var_to_aggregate, agg_function):
+def make_agg_var(agent, geog, geog_id, var_to_aggregate, agg_function, how_fillna = None):
     """
     Generator function for aggregation variables. Registers with orca.
     """
@@ -42,10 +42,14 @@ def func():
         # Fillna.
         # For certain functions, must add other options,
         # like puma value or neighboring value
-        if agg_function == 'sum':
-            series = series.fillna(0)
+        if how_fillna is not None:
+            series = how_fillna(series)
         else:
-            series = series.fillna(series.median())
+            if agg_function == 'sum':
+                series = series.fillna(0)
+            else:
+                series = series.fillna(method='ffill')
+                series = series.fillna(method='bfill')
 
         return series