rename argument to multiindex and add assertion

collect/cnrfc/cnrfc.py

@@ -240,7 +240,7 @@ def get_deterministic_forecast(cnrfc_id, truncate_historical=False, release=Fals
                                  'downloaded': dt.datetime.now().strftime('%Y-%m-%d %H:%M')}}
 
 
-def get_deterministic_forecast_watershed(watershed, date_string, acre_feet=False, pdt_convert=False, as_pdt=False, cnrfc_id=None, keep_stage=False):
+def get_deterministic_forecast_watershed(watershed, date_string, acre_feet=False, pdt_convert=False, as_pdt=False, cnrfc_id=None, multiindex=False):
     """
     from: https://www.cnrfc.noaa.gov/deterministicHourlyProductCSV.php
     https://www.cnrfc.noaa.gov/csv/2019040318_american_csv_export.zip
@@ -252,6 +252,7 @@ def get_deterministic_forecast_watershed(watershed, date_string, acre_feet=False
         pdt_convert (bool): 
         as_pdt (bool): 
         cnrfc_id (str): 
+        multiindex (bool): return df with MultiIndex (column headers with 2 rows)
     Returns:
         (dict): 
     """
@@ -303,8 +304,8 @@ def get_deterministic_forecast_watershed(watershed, date_string, acre_feet=False
     # convert kcfs to cfs; optional timezone conversions and optional conversion to acre-feet
     df, units = _apply_conversions(df, 'hourly', acre_feet, pdt_convert, as_pdt)
 
-    # Drop MultiIndex if not keep_stage (to preserve backwards compatibility)
-    if not keep_stage:
+    # Drop MultiIndex if not multindex (to preserve backwards compatibility)
+    if not multiindex:
         df.columns = df.columns.droplevel(level=1)
 
     # clean up
@@ -320,21 +321,40 @@ def get_deterministic_forecast_watershed(watershed, date_string, acre_feet=False
                                  'units': units,
                                  'downloaded': dt.datetime.now().strftime('%Y-%m-%d %H:%M')}}
 
-def get_deterministic_forecast_watershed_stage(watershed, date_string, pdt_convert=False, as_pdt=False, cnrfc_id=None):
+def get_deterministic_forecast_watershed_stage(watershed, date_string, pdt_convert=False, as_pdt=False, cnrfc_id=None, multiindex=False):
+    """
+    Return deterministic forecast for stage
+    from: https://www.cnrfc.noaa.gov/deterministicHourlyProductCSV.php
+    https://www.cnrfc.noaa.gov/csv/2019040318_american_csv_export.zip
+
+    Arguments:
+        watershed (str):
+        date_string (str):
+        acre_feet (bool): 
+        pdt_convert (bool): 
+        as_pdt (bool): 
+        cnrfc_id (str): 
+        multiindex (bool): return df with MultiIndex (column headers with 2 rows)
+    Returns:
+        (dict): 
+    """
 
     results = get_deterministic_forecast_watershed(watershed,
                                                    date_string,
                                                    acre_feet=False,
                                                    pdt_convert=pdt_convert,
                                                    as_pdt=as_pdt,
                                                    cnrfc_id=cnrfc_id,
-                                                   keep_stage=True)
+                                                   multiindex=True)
 
     info = results['info']
 
     df = results['data'].filter(regex='SSTG', axis=1)
 
-    df.columns = df.columns.droplevel(level=1)
+    assert not df.empty, 'no stage deterministic forecast'
+
+    if not multiindex:
+        df.columns = df.columns.droplevel(level=1)
 
     info.update({'units': 'ft'})