remove source data from local after patch run

Author: minhkhul

nssp/delphi_nssp/patch.py

@@ -41,7 +41,7 @@
 
 import sys
 from datetime import datetime, timedelta
-from os import listdir, makedirs, path, getcwd
+from os import listdir, makedirs, path
 from shutil import rmtree
 
 from delphi_utils import get_structured_logger, read_params
@@ -105,7 +105,7 @@ def patch():
 
     source_dir = params["patch"]["source_dir"]
     download_source = False
-    if not path.isdir(source_dir) or not listdir(source_dir): #no source dir or empty source dir
+    if not path.isdir(source_dir) or not listdir(source_dir):  # no source dir or empty source dir
         download_source = True
         get_source_data(params, logger)
     else:
@@ -145,8 +145,8 @@ def patch():
         run_module(params, logger)
         current_issue += timedelta(days=1)
 
-    # if download_source:
-    #     rmtree(source_dir)
+    if download_source:
+        rmtree(source_dir)
 
 
 if __name__ == "__main__":

nssp/delphi_nssp/pull.py

@@ -1,12 +1,12 @@
 # -*- coding: utf-8 -*-
 """Functions for pulling NSSP ER data."""
 
-import logging
 import functools
+import logging
 import sys
 import textwrap
-from typing import Optional
 from os import makedirs, path
+from typing import Optional
 
 import pandas as pd
 import paramiko
@@ -136,7 +136,13 @@ def pull_with_socrata_api(socrata_token: str, dataset_id: str):
     return results
 
 
-def pull_nssp_data(socrata_token: str, backup_dir: str, custom_run: bool, issue_date: Optional[str] = None, logger: Optional[logging.Logger] = None):
+def pull_nssp_data(
+    socrata_token: str,
+    backup_dir: str,
+    custom_run: bool,
+    issue_date: Optional[str] = None,
+    logger: Optional[logging.Logger] = None,
+):
     """Pull the NSSP ER visits primary dataset.
 
     https://data.cdc.gov/Public-Health-Surveillance/NSSP-Emergency-Department-Visit-Trajectories-by-St/rdmq-nq56/data_preview
@@ -166,7 +172,7 @@ def pull_nssp_data(socrata_token: str, backup_dir: str, custom_run: bool, issue_
             num_records=len(df_ervisits),
             source=source_filename,
         )
-    
+
     df_ervisits = df_ervisits.rename(columns={"week_end": "timestamp"})
     df_ervisits = df_ervisits.rename(columns=SIGNALS_MAP)
 
@@ -183,7 +189,11 @@ def pull_nssp_data(socrata_token: str, backup_dir: str, custom_run: bool, issue_
 
 
 def secondary_pull_nssp_data(
-    socrata_token: str, backup_dir: str, custom_run: bool, issue_date: Optional[str] = None, logger: Optional[logging.Logger] = None
+    socrata_token: str,
+    backup_dir: str,
+    custom_run: bool,
+    issue_date: Optional[str] = None,
+    logger: Optional[logging.Logger] = None,
 ):
     """Pull the latest NSSP ER visits secondary dataset.
 
@@ -207,9 +217,7 @@ def secondary_pull_nssp_data(
         socrata_results = pull_with_socrata_api(socrata_token, "7mra-9cq9")
         df_ervisits = pd.DataFrame.from_records(socrata_results)
         create_backup_csv(df_ervisits, backup_dir, custom_run, sensor="secondary", logger=logger)
-        logger.info("Number of records grabbed",
-                    num_records=len(df_ervisits),
-                    source="secondary Socrata API")
+        logger.info("Number of records grabbed", num_records=len(df_ervisits), source="secondary Socrata API")
 
     elif custom_run and logger.name == "delphi_nssp.patch":
         if issue_date is None:

nssp/delphi_nssp/run.py

@@ -152,7 +152,9 @@ def run_module(params, logger=None):
                 run_stats.append((max(dates), len(dates)))
 
     logger.info("Generating secondary signals")
-    secondary_df_pull = secondary_pull_nssp_data(socrata_token, backup_dir, custom_run=custom_run, issue_date=issue_date, logger=logger)
+    secondary_df_pull = secondary_pull_nssp_data(
+        socrata_token, backup_dir, custom_run=custom_run, issue_date=issue_date, logger=logger
+    )
     for signal in SECONDARY_SIGNALS:
         secondary_df_pull_signal = secondary_df_pull[secondary_df_pull["signal"] == signal]
         if secondary_df_pull_signal.empty:

nssp/setup.py

@@ -14,6 +14,7 @@
     "freezegun",
     "us",
     "paramiko",
+    "darker",
 ]
 
 setup(