From 98a27343e3a06f7ede11529b5c2eb1cdf4ca59ea Mon Sep 17 00:00:00 2001 From: Frederik Kratzert Date: Sun, 19 Oct 2025 10:31:52 +0000 Subject: [PATCH] Refactor: Rename method that loads cached sites file --- rivretrieve/australia.py | 6 +++--- rivretrieve/base.py | 8 ++------ rivretrieve/brazil.py | 6 +++--- rivretrieve/canada.py | 6 +++--- rivretrieve/chile.py | 6 +++--- rivretrieve/france.py | 6 +++--- rivretrieve/japan.py | 6 +++--- rivretrieve/poland.py | 6 +++--- rivretrieve/slovenia.py | 6 +++--- rivretrieve/southafrica.py | 6 +++--- rivretrieve/uk_ea.py | 6 +++--- rivretrieve/uk_nrfa.py | 4 ++-- rivretrieve/usa.py | 6 +++--- rivretrieve/utils.py | 7 +++---- 14 files changed, 40 insertions(+), 45 deletions(-) diff --git a/rivretrieve/australia.py b/rivretrieve/australia.py index e6da17e..5dc6f6a 100644 --- a/rivretrieve/australia.py +++ b/rivretrieve/australia.py @@ -19,9 +19,9 @@ class AustraliaFetcher(base.RiverDataFetcher): BOM_URL = "http://www.bom.gov.au/waterdata/services" @staticmethod - def get_gauge_ids() -> pd.DataFrame: - """Retrieves a DataFrame of available Australian gauge IDs.""" - return utils.load_sites_csv("australia") + def get_cached_metadata() -> pd.DataFrame: + """Retrieves a DataFrame of available Australian gauge IDs and metadata.""" + return utils.load_cached_metadata_csv("australia") @staticmethod def get_available_variables() -> tuple[str, ...]: diff --git a/rivretrieve/base.py b/rivretrieve/base.py index fd6746c..199a8ad 100644 --- a/rivretrieve/base.py +++ b/rivretrieve/base.py @@ -38,12 +38,8 @@ def get_data( @staticmethod @abc.abstractmethod - def get_gauge_ids() -> pd.DataFrame: - """Retrieves a DataFrame of available gauge IDs. - - Returns: - A pandas DataFrame containing gauge ID information. - """ + def get_cached_metadata() -> pd.DataFrame: + """Retrieves a DataFrame of available gauge IDs and metadata from a cached file.""" pass def get_metadata(self) -> pd.DataFrame: diff --git a/rivretrieve/brazil.py b/rivretrieve/brazil.py index b92da1f..9339182 100644 --- a/rivretrieve/brazil.py +++ b/rivretrieve/brazil.py @@ -42,9 +42,9 @@ def __init__(self, username: Optional[str] = None, password: Optional[str] = Non ) @staticmethod - def get_gauge_ids() -> pd.DataFrame: - """Retrieves a DataFrame of available Brazilian gauge IDs.""" - return utils.load_sites_csv("brazil") + def get_cached_metadata() -> pd.DataFrame: + """Retrieves a DataFrame of available Brazilian gauge IDs and metadata.""" + return utils.load_cached_metadata_csv("brazil") @staticmethod def get_available_variables() -> tuple[str, ...]: diff --git a/rivretrieve/canada.py b/rivretrieve/canada.py index 16599d2..3c0219b 100644 --- a/rivretrieve/canada.py +++ b/rivretrieve/canada.py @@ -27,9 +27,9 @@ class CanadaFetcher(base.RiverDataFetcher): HYDAT_PATH = DATA_DIR / "Hydat.sqlite3" @staticmethod - def get_gauge_ids() -> pd.DataFrame: - """Retrieves a DataFrame of available Canadian gauge IDs.""" - return utils.load_sites_csv("canada") + def get_cached_metadata() -> pd.DataFrame: + """Retrieves a DataFrame of available Canadian gauge IDs and metadata.""" + return utils.load_cached_metadata_csv("canada") @staticmethod def get_available_variables() -> tuple[str, ...]: diff --git a/rivretrieve/chile.py b/rivretrieve/chile.py index 4388fef..368835a 100644 --- a/rivretrieve/chile.py +++ b/rivretrieve/chile.py @@ -18,9 +18,9 @@ class ChileFetcher(base.RiverDataFetcher): """Fetches river gauge data from Chile's CR2 explorador.""" @staticmethod - def get_gauge_ids() -> pd.DataFrame: - """Retrieves a DataFrame of available Chilean gauge IDs.""" - return utils.load_sites_csv("chile") + def get_cached_metadata() -> pd.DataFrame: + """Retrieves a DataFrame of available Chilean gauge IDs and metadata.""" + return utils.load_cached_metadata_csv("chile") @staticmethod def get_available_variables() -> tuple[str, ...]: diff --git a/rivretrieve/france.py b/rivretrieve/france.py index dc0d47f..5aaadcd 100644 --- a/rivretrieve/france.py +++ b/rivretrieve/france.py @@ -17,9 +17,9 @@ class FranceFetcher(base.RiverDataFetcher): BASE_URL = "https://hubeau.eaufrance.fr/api/v2/hydrometrie/obs_elab" @staticmethod - def get_gauge_ids() -> pd.DataFrame: - """Retrieves a DataFrame of available French gauge IDs.""" - return utils.load_sites_csv("french") # Note: CSV file name is french_sites.csv + def get_cached_metadata() -> pd.DataFrame: + """Retrieves a DataFrame of available French gauge IDs and metadata.""" + return utils.load_cached_metadata_csv("french") # Note: CSV file name is french_sites.csv @staticmethod def get_available_variables() -> tuple[str, ...]: diff --git a/rivretrieve/japan.py b/rivretrieve/japan.py index 13fc1f9..ea80911 100644 --- a/rivretrieve/japan.py +++ b/rivretrieve/japan.py @@ -21,9 +21,9 @@ class JapanFetcher(base.RiverDataFetcher): BASE_URL = "http://www1.river.go.jp/cgi-bin/DspWaterData.exe" @staticmethod - def get_gauge_ids() -> pd.DataFrame: - """Retrieves a DataFrame of available Japanese gauge IDs.""" - return utils.load_sites_csv("japan") + def get_cached_metadata() -> pd.DataFrame: + """Retrieves a DataFrame of available Japanese gauge IDs and metadata.""" + return utils.load_cached_metadata_csv("japan") @staticmethod def get_available_variables() -> tuple[str, ...]: diff --git a/rivretrieve/poland.py b/rivretrieve/poland.py index 7f08532..2a6fc7d 100644 --- a/rivretrieve/poland.py +++ b/rivretrieve/poland.py @@ -25,9 +25,9 @@ class PolandFetcher(base.RiverDataFetcher): CACHE_FILE = Path(os.path.dirname(__file__)) / "data" / "poland.zarr" @staticmethod - def get_gauge_ids() -> pd.DataFrame: - """Retrieves a DataFrame of available Polish gauge IDs.""" - return utils.load_sites_csv("poland") + def get_cached_metadata() -> pd.DataFrame: + """Retrieves a DataFrame of available Polish gauge IDs and metadata.""" + return utils.load_cached_metadata_csv("poland") @staticmethod def get_available_variables() -> tuple[str, ...]: diff --git a/rivretrieve/slovenia.py b/rivretrieve/slovenia.py index f99dd81..6e6adc1 100644 --- a/rivretrieve/slovenia.py +++ b/rivretrieve/slovenia.py @@ -19,9 +19,9 @@ class SloveniaFetcher(base.RiverDataFetcher): BASE_URL = "https://vode.arso.gov.si/hidarhiv/pov_arhiv_tab.php" @staticmethod - def get_gauge_ids() -> pd.DataFrame: - """Retrieves a DataFrame of available Slovenian gauge IDs.""" - return utils.load_sites_csv("slovenia") + def get_cached_metadata() -> pd.DataFrame: + """Retrieves a DataFrame of available Slovenian gauge IDs and metadata.""" + return utils.load_cached_metadata_csv("slovenia") @staticmethod def get_available_variables() -> tuple[str, ...]: diff --git a/rivretrieve/southafrica.py b/rivretrieve/southafrica.py index 3973618..a28da4a 100644 --- a/rivretrieve/southafrica.py +++ b/rivretrieve/southafrica.py @@ -21,9 +21,9 @@ class SouthAfricaFetcher(base.RiverDataFetcher): BASE_URL = "https://www.dws.gov.za/Hydrology/Verified/HyData.aspx" @staticmethod - def get_gauge_ids() -> pd.DataFrame: - """Retrieves a DataFrame of available South African gauge IDs.""" - return utils.load_sites_csv("southAfrican") + def get_cached_metadata() -> pd.DataFrame: + """Retrieves a DataFrame of available South African gauge IDs and metadata.""" + return utils.load_cached_metadata_csv("southAfrican") @staticmethod def get_available_variables() -> tuple[str, ...]: diff --git a/rivretrieve/uk_ea.py b/rivretrieve/uk_ea.py index 986954f..0c67329 100644 --- a/rivretrieve/uk_ea.py +++ b/rivretrieve/uk_ea.py @@ -28,9 +28,9 @@ class UKEAFetcher(base.RiverDataFetcher): } @staticmethod - def get_gauge_ids() -> pd.DataFrame: - """Retrieves a DataFrame of available UK gauge IDs.""" - return utils.load_sites_csv("uk_ea") + def get_cached_metadata() -> pd.DataFrame: + """Retrieves a DataFrame of available UK gauge IDs and metadata.""" + return utils.load_cached_metadata_csv("uk_ea") @staticmethod def get_available_variables() -> tuple[str, ...]: diff --git a/rivretrieve/uk_nrfa.py b/rivretrieve/uk_nrfa.py index 9691a98..f6f6413 100644 --- a/rivretrieve/uk_nrfa.py +++ b/rivretrieve/uk_nrfa.py @@ -28,9 +28,9 @@ class UKNRFAFetcher(base.RiverDataFetcher): } @staticmethod - def get_gauge_ids() -> pd.DataFrame: + def get_cached_metadata() -> pd.DataFrame: """Retrieves a DataFrame of available NRFA gauge IDs from the cached CSV.""" - return utils.load_sites_csv("uk_nrfa") + return utils.load_cached_metadata_csv("uk_nrfa") def get_metadata(self) -> pd.DataFrame: """Fetches site metadata from the NRFA API and renames columns.""" diff --git a/rivretrieve/usa.py b/rivretrieve/usa.py index 733594a..d638e6c 100644 --- a/rivretrieve/usa.py +++ b/rivretrieve/usa.py @@ -15,9 +15,9 @@ class USAFetcher(base.RiverDataFetcher): """Fetches river gauge data from USGS NWIS.""" @staticmethod - def get_gauge_ids() -> pd.DataFrame: - """Retrieves a DataFrame of available USA gauge IDs.""" - return utils.load_sites_csv("usa") + def get_cached_metadata() -> pd.DataFrame: + """Retrieves a DataFrame of available USA gauge IDs and metadata.""" + return utils.load_cached_metadata_csv("usa") @staticmethod def get_available_variables() -> tuple[str, ...]: diff --git a/rivretrieve/utils.py b/rivretrieve/utils.py index 34198ec..890818b 100644 --- a/rivretrieve/utils.py +++ b/rivretrieve/utils.py @@ -14,8 +14,6 @@ logger = logging.getLogger(__name__) -DEFAULT_HEADERS = {"User-Agent": "Mozilla/5.0"} - def format_start_date(start_date: Optional[str]) -> str: """Formats the start date, defaulting to 1900-01-01 if None.""" @@ -60,12 +58,13 @@ def requests_retry_session( return session -def load_sites_csv(country_code: str) -> pd.DataFrame: +def load_cached_metadata_csv(country_code: str) -> pd.DataFrame: """Loads site data from a CSV file in the data directory.""" current_dir = os.path.dirname(__file__) file_path = os.path.join(current_dir, "cached_site_data", f"{country_code}_sites.csv") try: - return pd.read_csv(file_path, dtype={constants.GAUGE_ID: str}) + df = pd.read_csv(file_path, dtype={constants.GAUGE_ID: str}) + return df.set_index(constants.GAUGE_ID) except FileNotFoundError: logger.error(f"Site file not found: {file_path}") raise