Make linter happy

Author: dshemetov

usafacts/delphi_usafacts/run.py

@@ -65,6 +65,27 @@
 ]
 
 
+def add_nancodes(df, smoother):
+    """Add nancodes to the dataframe."""
+    idx = pd.IndexSlice
+
+    # Default nancodes
+    df["missing_val"] = Nans.NOT_MISSING
+    df["missing_se"] = Nans.NOT_APPLICABLE
+    df["missing_sample_size"] = Nans.NOT_APPLICABLE
+
+    # Mark early smoothing entries as data insufficient
+    if smoother == "seven_day_average":
+        df.sort_index(inplace=True)
+        min_time_value = df.index.min()[0] + 6 * pd.Timedelta(days=1)
+        df.loc[idx[:min_time_value, :], "missing_val"] = Nans.DATA_INSUFFICIENT
+
+    # Mark any remaining nans with unknown
+    remaining_nans_mask = df["val"].isnull() & (df["missing_val"] == Nans.NOT_MISSING)
+    df.loc[remaining_nans_mask, "missing_val"] = Nans.UNKNOWN
+    return df
+
+
 def run_module(params: Dict[str, Dict[str, Any]]):
     """Run the usafacts indicator.
 
@@ -114,30 +135,16 @@ def run_module(params: Dict[str, Dict[str, Any]]):
         # Aggregate to appropriate geographic resolution
         df = geo_map(df, geo_res, sensor)
         df.set_index(["timestamp", "geo_id"], inplace=True)
-        idx = pd.IndexSlice
 
         # Smooth
-        smooth_obj, smoother_prefix, is_smooth_wip, smoother_lag = SMOOTHERS_MAP[smoother]
+        smooth_obj, smoother_prefix, _, smoother_lag = SMOOTHERS_MAP[smoother]
         df["val"] = df[sensor].groupby(level=1).transform(smooth_obj.smooth)
 
         # USAFacts is not a survey indicator
         df["se"] = np.nan
         df["sample_size"] = np.nan
 
-        # Default missing code
-        df["missing_val"] = Nans.NOT_MISSING
-        df["missing_se"] = Nans.NOT_APPLICABLE
-        df["missing_sample_size"] = Nans.NOT_APPLICABLE
-
-        # Mark early smoothing entries as data insufficient
-        if smoother == "seven_day_average":
-            df.sort_index(inplace=True)
-            min_time_value = df.index.min()[0] + 6 * pd.Timedelta(days=1)
-            df.loc[idx[:min_time_value, :], "missing_val"] = Nans.DATA_INSUFFICIENT
-
-        # Mark any remaining nans with unknown
-        remaining_nans_mask = df["val"].isnull() & (df["missing_val"] == Nans.NOT_MISSING)
-        df.loc[remaining_nans_mask, "missing_val"] = Nans.UNKNOWN
+        df = add_nancodes(df, smoother)
 
         df.reset_index(inplace=True)
         sensor_name = SENSOR_NAME_MAP[sensor][0]