ref: split extract_useful_data fn

services/grib/service.py

@@ -123,26 +123,8 @@ class Grib:
         )
 
     @staticmethod
-    def extract_useful_data(target_time: str = None):
-        """Download and process GRIB files into csv of requested parameters
-
-        :param target_time: well formed ISO 8601 time string, we will try to download GRIB available just before it
-
-        :returns: filenames of the resulting csv
-        """
-        if not target_time:
-            target_time = Grib.fresh_grib_time()
-        else:
-            target_time = Grib.fresh_grib_time(datetime.fromisoformat(target_time))
-        save_to = f"{target_time.strftime('%Y%m%d')}-{str((target_time.hour // 6) * 6).zfill(2)}.csv"
-
-        for forecast_hour in range(MAX_FORECAST_HOUR):
-            # download gribs
-            ds_atmos_file, ds_wave_file = Grib.download_fresh_grib(
-                target_time=target_time, prod_hour=forecast_hour
-            )
-
-            # filter atmos to requested variables
+    def grib_to_dataframe(ds_atmos_file, ds_wave_file, forecast_hour):
+        """Work with open GRIB file, extracting data into pandas dataframe"""
         with xr.open_dataset(ds_atmos_file, engine="pynio") as ds_atmos:
             filtered_ds_atmos = ds_atmos.get(ATMOS_PARAM_NAMES) or ds_atmos.get(
                 [p for p in ATMOS_PARAM_NAMES if not p == "APCP_P8_L1_GLL0_acc"]
@@ -172,9 +154,7 @@ class Grib:
 
                 # concatinate atmos and wave into a single dataset
                 combined_product = filtered_ds_atmos.merge(
-                        filtered_ds_wave.reindex_like(
-                            filtered_ds_atmos, method="nearest"
-                        )
+                    filtered_ds_wave.reindex_like(filtered_ds_atmos, method="nearest")
                 )
 
                 # transfer to pandas
@@ -189,7 +169,11 @@ class Grib:
         remapped_longitudes = longitudes.map(map_function)
         df["longitude"] = remapped_longitudes
         df["latitude"] = latitudes
-            # dump datafrate to csv on disk
+        return df
+
+    @staticmethod
+    def dump_df_to_csv(df, forecast_hour, save_to):
+        """Dump pandas dataframe to CSV on disk"""
         if forecast_hour == 0:
             df.to_csv(
                 os.path.join(
@@ -206,6 +190,31 @@ class Grib:
                 header=False,
             )
 
+    @staticmethod
+    def extract_useful_data(target_time: str = None):
+        """Download and process GRIB files into csv of requested parameters
+
+        :param target_time: well formed ISO 8601 time string, we will try to download GRIB available just before it
+
+        :returns: filenames of the resulting csv
+        """
+        if not target_time:
+            target_time = Grib.fresh_grib_time()
+        else:
+            target_time = Grib.fresh_grib_time(datetime.fromisoformat(target_time))
+        save_to = f"{target_time.strftime('%Y%m%d')}-{str((target_time.hour // 6) * 6).zfill(2)}.csv"
+
+        for forecast_hour in range(MAX_FORECAST_HOUR):
+            # download gribs
+            ds_atmos_file, ds_wave_file = Grib.download_fresh_grib(
+                target_time=target_time, prod_hour=forecast_hour
+            )
+
+            # filter atmos to requested variables
+            df = Grib.grib_to_dataframe(ds_atmos_file, ds_wave_file, forecast_hour)
+            # dump datafrate to csv on disk
+            Grib.dump_df_to_csv(df, forecast_hour, save_to)
+
             # clean up grib files
             os.remove(ds_wave_file)
             os.remove(ds_atmos_file)