From 2092571b0192ba681442a044a363f3f4ee043eda Mon Sep 17 00:00:00 2001 From: Daniel Howe Date: Thu, 14 Mar 2019 16:56:21 +1100 Subject: [PATCH] Add strict matching for basin names --- waternsw_grabber/waternsw_grabber.py | 16 +++++----------- 1 file changed, 5 insertions(+), 11 deletions(-) diff --git a/waternsw_grabber/waternsw_grabber.py b/waternsw_grabber/waternsw_grabber.py index 3676bf6..5a65df2 100644 --- a/waternsw_grabber/waternsw_grabber.py +++ b/waternsw_grabber/waternsw_grabber.py @@ -230,14 +230,14 @@ def open_browser(download_dir): return driver -def telemetered_bore_downloader(basin, +def telemetered_bore_downloader(basin_name, download_dir=None, start_date=None, end_date=None): """Download multiple records from telemetered bore. Args: - basin: basin name or code (string) + basin_name: basin name (string) download_dir: path to where downloaded files will be saved start_date: start date (string YYYY-MM-DD format) end_date: end date (string YYYY-MM-DD format) @@ -250,15 +250,9 @@ def telemetered_bore_downloader(basin, basins = get_basins() # Check if full basin name was provided - if basin in basins['Basin name'].values: - basin_name = basin - else: - try: - # Check ift two-letter basin code was provided - basin_name = basins.groupby('Basin code').first().loc[basin][ - 'Basin name'] - except KeyError: - raise KeyError('Basin {} not found'.format(basin)) + basin_names = basins['Basin name'].unique() + if basin_name not in basin_names: + raise ValueError("'basin_name' must be one of:\n" + '\n'.join(basin_names)) # Get list of bore IDs from selected basin bore_ids = basins[basins['Basin name'] == basin_name].index.values