diff --git a/src/cosecha/__init__.py b/src/cosecha/__init__.py index a0ca0a7..d985f58 100644 --- a/src/cosecha/__init__.py +++ b/src/cosecha/__init__.py @@ -7,6 +7,7 @@ from cosecha import exceptions from cosecha._logging import configure_logger from cosecha.reaping import ( + ASOSReaper, GriddedReaper, MRMSReaper, NWPReaper, @@ -20,6 +21,7 @@ __version__ = "999" __all__ = [ + "ASOSReaper", "GriddedReaper", "MRMSReaper", "NWPReaper", diff --git a/src/cosecha/exceptions.py b/src/cosecha/exceptions.py index 5617239..a35c1c9 100644 --- a/src/cosecha/exceptions.py +++ b/src/cosecha/exceptions.py @@ -4,6 +4,7 @@ __all__ = [ "APIError", + "DataNotFoundError", "DateRangeError", "InvalidSiteError", "ReaperError", @@ -27,5 +28,9 @@ class APIError(ReaperError): """Raised when an API call fails.""" +class DataNotFoundError(ReaperError): + """Raised when a query returns no data.""" + + class TransformationError(ReaperError): """Raised when data transformation fails.""" diff --git a/src/cosecha/reaping/__init__.py b/src/cosecha/reaping/__init__.py index 29cef37..b3193c9 100644 --- a/src/cosecha/reaping/__init__.py +++ b/src/cosecha/reaping/__init__.py @@ -6,6 +6,7 @@ from __future__ import annotations +from cosecha.reaping.asos import ASOSReaper from cosecha.reaping.base import GriddedReaper, TimeSeriesReaper from cosecha.reaping.mrms import MRMSReaper from cosecha.reaping.nwis import USGSNWISReaper @@ -13,6 +14,7 @@ from cosecha.reaping.usace import ReservoirReaper __all__ = [ + "ASOSReaper", "GriddedReaper", "MRMSReaper", "NWPReaper", diff --git a/src/cosecha/reaping/asos.py b/src/cosecha/reaping/asos.py new file mode 100644 index 0000000..1804d95 --- /dev/null +++ b/src/cosecha/reaping/asos.py @@ -0,0 +1,169 @@ +"""IEM ASOS data reapers. + +This module provides implementations for harvesting ASOS observations +from the Iowa Environmental Mesonet (IEM) API. +""" + +from __future__ import annotations + +from io import StringIO +from typing import Any +from urllib.parse import urlencode + +import pandas as pd +import tiny_retriever + +from cosecha._logging import logger +from cosecha._utils import apply_ts_transformations, wrap_errors +from cosecha.exceptions import APIError, DataNotFoundError, DateRangeError +from cosecha.reaping.base import TimeSeriesReaper + +__all__ = [ + "ASOSReaper", +] + +BASE_URL = "https://mesonet.agron.iastate.edu/cgi-bin/request/asos.py" +_IEM_ALL_VARS = "all" + + +class ASOSReaper(TimeSeriesReaper): + """Reaper for IEM ASOS data.""" + + def _validate_params(self) -> None: + """Validate initialization parameters. + + Raises + ------ + DateRangeError + If dates are invalid. + """ + if self.start_date > self.end_date: + raise DateRangeError( + f"start_date ({self.start_date}) must be <= end_date ({self.end_date})" + ) + + def __init__( + self, + start_date: str, + end_date: str, + state: str | None = None, + variable: str | list[str] | None = None, + report_type: int | list[int] | None = None, + transformations: dict[str, Any] | None = None, + timeout: int = 120, + ) -> None: + """Fetch data from IEM ASOS API. + + Parameters + ---------- + start_date : str + Start date in ISO 8601 format (YYYY-MM-DDTHH:MMZ). + end_date : str + End date in ISO 8601 format (YYYY-MM-DDTHH:MMZ). + state : str | None, optional + State abbreviation (e.g., 'TX'). If None (default), fetches from + all networks (IEM limits this to a 24-hour window). + variable : str | list[str] | None, optional + Variable to fetch (e.g., 'p01i', 'tmpf', or ['p01i', 'tmpf']). + If None, fetches 'all'. + report_type : int | list[int] | None, optional + IEM report type filter: 1 (HFMETAR/5-min), 3 (routine), 4 (specials). + If None (default), returns all report types. + transformations : dict[str, Any], optional + Optional transformations to apply to the data. + timeout : int, optional + Request timeout in seconds, by default 120. + """ + super().__init__() + self.state = state + self.network = f"{state.upper()}_ASOS" if state else None + + try: + self.start_date = pd.to_datetime(start_date) + self.end_date = pd.to_datetime(end_date) + except Exception as e: + raise DateRangeError(f"Could not parse date: {e}") from e + + if variable is None: + self.data_vars = [_IEM_ALL_VARS] + elif isinstance(variable, str): + self.data_vars = [variable] + else: + self.data_vars = variable + + self.transformations = transformations + self.timeout = timeout + + if report_type is None: + self.report_type = None + elif isinstance(report_type, int): + self.report_type = [report_type] + else: + self.report_type = report_type + + self._validate_params() + logger.debug( + f"Initialized {self.__class__.__name__}: " + f"network={self.network}, dates={self.start_date} to {self.end_date}, " + f"data={self.data_vars}" + ) + + def _build_url(self) -> str: + """Build the IEM ASOS request URL with query parameters.""" + params = { + "year1": self.start_date.year, + "month1": self.start_date.month, + "day1": self.start_date.day, + "hour1": self.start_date.hour, + "minute1": self.start_date.minute, + "year2": self.end_date.year, + "month2": self.end_date.month, + "day2": self.end_date.day, + "hour2": self.end_date.hour, + "minute2": self.end_date.minute, + "tz": "Etc/UTC", + "format": "comma", + "latlon": "yes", + "data": self.data_vars, + } + if self.network is not None: + params["network"] = self.network + if self.report_type is not None: + params["report_type"] = self.report_type + return f"{BASE_URL}?{urlencode(params, doseq=True)}" + + def _fetch(self, url: str) -> str: + """Fetch CSV text from IEM via tiny_retriever.""" + with wrap_errors( + APIError, f"Failed to fetch ASOS data for {self.network or 'all networks'}" + ): + return tiny_retriever.fetch(url, "text", timeout=self.timeout) + + def _parse_response(self, text: str) -> pd.DataFrame: + """Parse IEM CSV text into a DataFrame. + + The first 5 rows of IEM ASOS output are comments (skiprows=5). + """ + df = pd.read_csv(StringIO(text), skiprows=5) + if df.empty: + raise DataNotFoundError( + f"ASOS returned no data for network {self.network or 'all networks'} and time range {self.start_date} to {self.end_date}" + ) + logger.debug(f"Fetched {len(df)} records from ASOS for {self.network or 'all networks'}") + return df + + def _reap(self) -> pd.DataFrame: + """Fetch data from ASOS and return as a pandas DataFrame.""" + logger.info( + f"Reaping ASOS data: network={self.network or 'all networks'}, data={self.data_vars}" + ) + + url = self._build_url() + text = self._fetch(url) + df = self._parse_response(text) + + if self.transformations and not df.empty: + df = apply_ts_transformations(df, self.transformations) + + logger.info(f"Reaped {len(df)} records from {self.network or 'all networks'}") + return df diff --git a/tests/test_reaping/test_asos_reapers.py b/tests/test_reaping/test_asos_reapers.py new file mode 100644 index 0000000..956536c --- /dev/null +++ b/tests/test_reaping/test_asos_reapers.py @@ -0,0 +1,162 @@ +"""Tests for ASOS reapers.""" + +from __future__ import annotations + +from unittest.mock import patch + +import pandas as pd +import pytest + +from cosecha.exceptions import APIError, DataNotFoundError, DateRangeError +from cosecha.reaping.asos import ASOSReaper + + +class TestASOSReaper: + """Tests for ASOSReaper.""" + + def test_initialization_valid(self): + """Test valid initialization.""" + reaper = ASOSReaper( + state="TX", + variable="p01i", + start_date="2026-04-12", + end_date="2026-04-13", + ) + assert reaper.state == "TX" + assert reaper.network == "TX_ASOS" + assert reaper.start_date == pd.Timestamp("2026-04-12") + assert reaper.end_date == pd.Timestamp("2026-04-13") + assert reaper.data_vars == ["p01i"] + + def test_initialization_multiple_data_vars(self): + """Test initialization with multiple data variables.""" + reaper = ASOSReaper( + state="IA", + variable=["p01i", "tmpf"], + start_date="2026-04-12", + end_date="2026-04-13", + ) + assert reaper.data_vars == ["p01i", "tmpf"] + assert reaper.network == "IA_ASOS" + + def test_initialization_all_network(self): + """Test initialization with no state fetches all networks.""" + reaper = ASOSReaper( + start_date="2026-04-12", + end_date="2026-04-13", + ) + assert reaper.network is None + assert reaper.data_vars == ["all"] + + def test_invalid_date_range(self): + """Test initialization fails with invalid date range.""" + with pytest.raises(DateRangeError): + ASOSReaper( + state="TX", + variable="p01i", + start_date="2026-04-13", + end_date="2026-04-12", + ) + + def test_invalid_start_date(self): + """Test initialization fails with unparsable start date.""" + with pytest.raises(DateRangeError, match="Could not parse date"): + ASOSReaper( + state="TX", + variable="p01i", + start_date="not-a-date", + end_date="2026-04-13", + ) + + @pytest.mark.network + def test_reap_network(self): + """Test live network retrieval from ASOS.""" + reaper = ASOSReaper( + state="TX", + variable="p01i", + start_date="2022-01-01", + end_date="2022-01-02", + ) + harvested = reaper.reap() + assert len(harvested) > 0 + assert "p01i" in harvested.columns + assert isinstance(harvested, pd.DataFrame) + + @patch("cosecha.reaping.asos.tiny_retriever.fetch") + def test_reap_success(self, mock_fetch): + """Test successful data retrieval.""" + mock_fetch.return_value = ( + "# comment 1\n" + "# comment 2\n" + "# comment 3\n" + "# comment 4\n" + "# comment 5\n" + "station,valid,lon,lat,p01i\n" + "AUS,2026-04-12 00:00,-97.6698,30.1945,0.01\n" + "AUS,2026-04-12 01:00,-97.6698,30.1945,0.02\n" + ) + + reaper = ASOSReaper( + state="TX", + variable="p01i", + start_date="2026-04-12", + end_date="2026-04-12", + ) + + harvested = reaper.reap() + + assert len(harvested) == 2 + assert isinstance(harvested, pd.DataFrame) + assert "p01i" in harvested.columns + assert "station" in harvested.columns + mock_fetch.assert_called_once() + + url = mock_fetch.call_args[0][0] + assert "network=TX_ASOS" in url + assert "data=p01i" in url + + @patch("cosecha.reaping.asos.tiny_retriever.fetch") + def test_reap_api_error(self, mock_fetch): + """Test error handling when API fails.""" + mock_fetch.side_effect = Exception("Connection failed") + + reaper = ASOSReaper( + state="TX", + start_date="2026-04-12", + end_date="2026-04-12", + ) + + with pytest.raises(APIError, match="Failed to fetch ASOS data for TX_ASOS"): + reaper.reap() + + mock_fetch.assert_called_once() + + @patch("cosecha.reaping.asos.tiny_retriever.fetch") + def test_reap_empty_response(self, mock_fetch): + """Test reap raises DataNotFoundError on empty output.""" + mock_fetch.return_value = "# 1\n# 2\n# 3\n# 4\n# 5\nstation,valid,lon,lat,p01i\n" + + reaper = ASOSReaper( + state="TX", + start_date="2026-04-12", + end_date="2026-04-12", + ) + + with pytest.raises(DataNotFoundError, match="no data"): + reaper.reap() + + @patch("cosecha.reaping.asos.tiny_retriever.fetch") + def test_reap_with_transformations(self, mock_fetch): + """Test reap applies format transformations when not empty.""" + mock_fetch.return_value = "# 1\n# 2\n# 3\n# 4\n# 5\nstation,p01i\nAUS,0.01\n" + + reaper = ASOSReaper( + state="TX", + start_date="2026-04-12", + end_date="2026-04-12", + transformations={"rename_columns": {"p01i": "precip"}}, + ) + + harvested = reaper.reap() + assert "precip" in harvested.columns + assert "p01i" not in harvested.columns