migrate safegraph.run back onto using functools.partial

Author: sgsmob

safegraph/delphi_safegraph/process.py

@@ -15,6 +15,7 @@
 # Base file name for raw data CSVs.
 CSV_NAME = 'social-distancing.csv.gz'
 
+
 def validate(df):
     """Confirms that a data frame has only one date."""
     timestamps = df['date_range_start'].apply(date_from_timestamp)
@@ -240,8 +241,7 @@ def process_window(df_list: List[pd.DataFrame],
                              index=False, )
 
 
-def process(current_filename: str,
-            previous_filenames: List[str],
+def process(filenames: List[str],
             signal_names: List[str],
             wip_signal,
             geo_resolutions: List[str],
@@ -250,11 +250,11 @@ def process(current_filename: str,
     as averaged over the previous week.
     Parameters
     ----------
-    current_filename: str
-        path to file holding the target date's data.
-    previous_filenames: List[str]
-        paths to files holding data from each day in the week preceding the
-        target date.
+    current_filename: List[str]
+        paths to files holding data.
+        The first entry of the list should correspond to the target date while
+        the remaining entries should correspond to the dates from each day in
+        the week preceding the target date.
     signal_names: List[str]
         signal names to be processed for a single date.
         A second version of each such signal named {SIGNAL}_7d_avg will be
@@ -274,8 +274,8 @@ def process(current_filename: str,
     one for the data averaged over the previous week to
     {export_dir}/{date}_{resolution}_{signal}_7d_avg.csv.
     """
-    past_week = [pd.read_csv(current_filename)]
-    for fname in previous_filenames:
+    past_week = []
+    for fname in filenames:
         if os.path.exists(fname):
             past_week.append(pd.read_csv(fname))
 
@@ -286,8 +286,8 @@ def process(current_filename: str,
                    export_dir)
     # ...then as part of the whole window.
     process_window(past_week,
-                  add_prefix(add_suffix(signal_names, '_7d_avg'),
-                             wip_signal,
-                             'wip_'),
-                  geo_resolutions,
-                  export_dir)
+                   add_prefix(add_suffix(signal_names, '_7d_avg'),
+                              wip_signal,
+                              'wip_'),
+                   geo_resolutions,
+                   export_dir)

safegraph/delphi_safegraph/run.py

@@ -3,6 +3,7 @@
 when the module is run with `python -m MODULE_NAME`.
 """
 import glob
+import functools
 import multiprocessing as mp
 import subprocess
 
@@ -24,22 +25,13 @@ def run_module():
     aws_endpoint = params["aws_endpoint"]
     wip_signal = params["wip_signal"]
 
-    def process_file(current_filename):
-        """Wrapper around `process()` that only takes a single argument.
-
-        A single argument function is necessary to use `pool.map()` below.
-        Because each call to `process()` has two arguments that are dependent
-        on the input file name (`current_filename` and `previous_filenames`),
-        we choose to use this wrapper rather than something like
-        `functools.partial()`.
-        """
-        return process(current_filename,
-                       files_in_past_week(current_filename),
-                       signal_names=SIGNALS,
-                       wip_signal=wip_signal,
-                       geo_resolutions=GEO_RESOLUTIONS,
-                       export_dir=export_dir,
-                       )
+    single_arg_process = functools.partial(
+        process,
+        signal_names=SIGNALS,
+        wip_signal=wip_signal,
+        geo_resolutions=GEO_RESOLUTIONS,
+        export_dir=export_dir,
+    )
 
     # Update raw data
     # Why call subprocess rather than using a native Python client, e.g. boto3?
@@ -60,5 +52,11 @@ def process_file(current_filename):
     files = glob.glob(f'{raw_data_dir}/social-distancing/**/*.csv.gz',
                       recursive=True)
 
+    files_with_previous_weeks = []
+    for fname in files:
+        previous_week = [fname]
+        previous_week.extend(files_in_past_week(fname))
+        files_with_previous_weeks.append(previous_week)
+
     with mp.Pool(n_core) as pool:
-        pool.map(process_file, files)
+        pool.map(single_arg_process, files_with_previous_weeks)

safegraph/tests/test_process.py

@@ -137,9 +137,9 @@ def test_process(self, tmp_path):
         export_dir = tmp_path / 'export'
         export_dir.mkdir()
 
-        process('raw_data/small_raw_data_0.csv',
-                # File 2 does not exist.
-                ['raw_data/small_raw_data_1.csv',
+        process(['raw_data/small_raw_data_0.csv',
+                 'raw_data/small_raw_data_1.csv',
+                 # File 2 does not exist.
                  'raw_data/small_raw_data_2.csv',
                  'raw_data/small_raw_data_3.csv', ],
                 SIGNALS,