diff --git a/waternsw_grabber/waternsw_grabber.py b/waternsw_grabber/waternsw_grabber.py index 9c2b443..19a4fa6 100644 --- a/waternsw_grabber/waternsw_grabber.py +++ b/waternsw_grabber/waternsw_grabber.py @@ -250,11 +250,11 @@ def extract_definitions(input_dir, output_dir): output_dir: path to save csv files """ - # Get telemetered site data + # Get basin info for telemetered site data csv_name = os.path.join( os.path.dirname(os.path.dirname(__file__)), 'data', 'telemetered-sites.csv') - master = pd.read_csv(csv_name, index_col=0) + basins = pd.read_csv(csv_name, index_col=0) # Find zip files zip_names = [f for f in os.listdir(input_dir) if f.endswith('.zip')] @@ -318,8 +318,8 @@ def extract_definitions(input_dir, output_dir): sites = sites.set_index('ID') # Get basin from master site dataframe - sites['Basin name'] = master.loc[sites.index, 'Basin name'] - sites['Basin code'] = master.loc[sites.index, 'Basin code'] + sites['Basin name'] = basins.loc[sites.index, 'Basin name'] + sites['Basin code'] = basins.loc[sites.index, 'Basin code'] # Save variable definitions variables = pd.DataFrame( @@ -373,6 +373,12 @@ def extract_records(input_dir, output_dir, clean_up=False): clean_up: delete original zip archive after extracting it """ + # Get basin info for telemetered site data + csv_name = os.path.join( + os.path.dirname(os.path.dirname(__file__)), 'data', + 'telemetered-sites.csv') + basins = pd.read_csv(csv_name, index_col=0) + # Find zip files zip_names = [f for f in os.listdir(input_dir) if f.endswith('.zip')]