diff --git a/rivretrieve/base.py b/rivretrieve/base.py index 199a8ad..5c8ec52 100644 --- a/rivretrieve/base.py +++ b/rivretrieve/base.py @@ -51,7 +51,7 @@ def get_metadata(self) -> pd.DataFrame: """ # Default implementation returns an empty DataFrame. # Subclasses should override this method if metadata is available. - return pd.DataFrame().set_index("gauge_id") + raise NotImplementedError @staticmethod @abc.abstractmethod diff --git a/rivretrieve/poland.py b/rivretrieve/poland.py index 2a6fc7d..b33894a 100644 --- a/rivretrieve/poland.py +++ b/rivretrieve/poland.py @@ -1,5 +1,6 @@ """Fetcher for Polish river gauge data from IMGW.""" +import io import logging import os import re @@ -23,10 +24,48 @@ class PolandFetcher(base.RiverDataFetcher): BASE_URL = "https://danepubliczne.imgw.pl/data/dane_pomiarowo_obserwacyjne/dane_hydrologiczne/" CACHE_FILE = Path(os.path.dirname(__file__)) / "data" / "poland.zarr" + METADATA_URL = ( + "https://danepubliczne.imgw.pl/data/dane_pomiarowo_obserwacyjne/dane_hydrologiczne/lista_stacji_hydro.csv" + ) + METADATA_CSV = Path(os.path.dirname(__file__)) / "cached_site_data" / "poland_sites.csv" + + @staticmethod + def get_metadata(): + """Downloads the metadata CSV file and converts it into a pandas DataFrame.""" + logger.info(f"Downloading metadata from {PolandFetcher.METADATA_URL}") + try: + r = utils.requests_retry_session().get(PolandFetcher.METADATA_URL) + r.raise_for_status() + + # The file is encoded in cp1250 + df = pd.read_csv(io.StringIO(r.content.decode("cp1250")), header=None, dtype=str) + + logger.info("Successfully downloaded and read metadata.") + + # Assign column names based on manual inspection of data on the https:// server + col_names = [ + constants.GAUGE_ID, + constants.STATION_NAME, + constants.RIVER, + "Kod Hydro", # Seems to be an alternativ station id. + ] + df.columns = col_names + + # Strip potential whitespace from gauge IDs + df[constants.GAUGE_ID] = df[constants.GAUGE_ID].str.strip() + + return df.set_index(constants.GAUGE_ID) + + except requests.exceptions.RequestException as e: + logger.error(f"Error downloading metadata: {e}") + raise + except Exception as e: + logger.error(f"Error processing or saving metadata: {e}") + raise @staticmethod def get_cached_metadata() -> pd.DataFrame: - """Retrieves a DataFrame of available Polish gauge IDs and metadata.""" + """Loads cache metadata.""" return utils.load_cached_metadata_csv("poland") @staticmethod diff --git a/tests/test_data/poland_metadata.csv b/tests/test_data/poland_metadata.csv new file mode 100644 index 0000000..1f817af --- /dev/null +++ b/tests/test_data/poland_metadata.csv @@ -0,0 +1,3 @@ +" 152140010","BIELINEK","Odra (1)","00101" +" 153140020","WIDUCHOWA","Odra (1)","00110" +" 153140030","GRYFINO","Odra (1)","00111" \ No newline at end of file diff --git a/tests/test_poland.py b/tests/test_poland.py index 1dd2061..6ad5983 100644 --- a/tests/test_poland.py +++ b/tests/test_poland.py @@ -136,6 +136,31 @@ def mock_get_side_effect(url, *args, **kwargs): self.assertEqual(parsed_df[constants.TIME_INDEX].min(), pd.to_datetime("2022-01-01")) self.assertEqual(parsed_df[constants.TIME_INDEX].max(), pd.to_datetime("2022-02-28")) + @patch("rivretrieve.utils.requests_retry_session") + def test_get_metadata(self, mock_requests_session): + mock_session = MagicMock() + mock_requests_session.return_value = mock_session + + with open(self.test_data_dir / "poland_metadata.csv", "rb") as f: + mock_content = f.read() + + mock_response = MagicMock() + mock_response.content = mock_content + mock_response.raise_for_status = MagicMock() + mock_session.get.return_value = mock_response + + metadata_df = self.fetcher.get_metadata() + + expected_data = { + constants.GAUGE_ID: ["152140010", "153140020", "153140030"], + constants.STATION_NAME: ["BIELINEK", "WIDUCHOWA", "GRYFINO"], + constants.RIVER: ["Odra (1)", "Odra (1)", "Odra (1)"], + "Kod Hydro": ["00101", "00110", "00111"], + } + expected_df = pd.DataFrame(expected_data).set_index(constants.GAUGE_ID) + + assert_frame_equal(metadata_df, expected_df) + if __name__ == "__main__": unittest.main()