diff --git a/waternsw_grabber/waternsw_grabber.py b/waternsw_grabber/waternsw_grabber.py index ecc1073..0575048 100644 --- a/waternsw_grabber/waternsw_grabber.py +++ b/waternsw_grabber/waternsw_grabber.py @@ -234,12 +234,13 @@ def telemetered_bore_downloader(bore_ids, start_date, end_date, download_dir): driver.quit() -def extract_records(input_dir, output_dir): +def extract_records(input_dir, output_dir, clean_up=False): """Extract downloaded bore records. Args: input_dir: path to downloaded zip archives output_dir: path to save csv files + clean_up: delete original zip archive after extracting it """ # Find zip files @@ -311,3 +312,7 @@ def extract_records(input_dir, output_dir): # Export to csv df.to_csv(csv_name, float_format='%0.3f') + + if clean_up: + # Remove original zip archive + os.remove(os.path.join(input_dir, zip_name))