From 37885441b3cd0ff54f27d5580e7d7639897fc977 Mon Sep 17 00:00:00 2001 From: Elise Hinman Date: Wed, 26 Aug 2026 12:06:26 -0500 Subject: [PATCH 1/5] simplify wqp module and combine get_result with what_ functions. created by Kiro --- dataretrieval/wqp.py | 94 +++++++++++++++++++++++++++++++++----------- 1 file changed, 71 insertions(+), 23 deletions(-) diff --git a/dataretrieval/wqp.py b/dataretrieval/wqp.py index 8b9fbb8d..06aba0fa 100644 --- a/dataretrieval/wqp.py +++ b/dataretrieval/wqp.py @@ -75,6 +75,52 @@ "Station", ] +# Valid dataProfile values per service, keyed by (service, legacy: bool). +# None means the service does not accept a dataProfile parameter at all. +_PROFILES: dict[tuple[str, bool], list[str] | None] = { + ("Result", True): result_profiles_legacy, + ("Result", False): result_profiles_wqx3, + ("Activity", True): activity_profiles_legacy, + ("Activity", False): None, + ("Station", True): None, + ("Station", False): None, + ("Organization", True): None, + ("Project", True): None, + ("ActivityMetric", True): None, + ("BiologicalMetric", True): None, + ("ProjectMonitoringLocationWeighting", True): None, + ("ResultDetectionQuantitationLimit", True): None, +} + + +def _validate_profile(service: str, legacy: bool, kwargs: dict[str, Any]) -> None: + """Check that ``dataProfile`` is valid for *service*, or absent when unsupported. + + Raises ``ValueError`` if the service does not accept profiles and one was + passed, or if the profile is not in the service's known list. When the + service is Result with ``legacy=False`` and no profile is given, defaults + to ``'fullPhysChem'`` (mutates *kwargs* in place). + """ + profile = kwargs.get("dataProfile") + key = (service, legacy) + valid = _PROFILES.get(key) + + if valid is None and profile is not None: + raise ValueError( + f"The {service!r} service does not accept a dataProfile parameter. " + f"Remove dataProfile={profile!r} from your call." + ) + + if valid is not None and profile is not None: + kind = "legacy" if legacy else "WQX3.0" + require_one_of( + profile, valid, name="dataProfile", context=f"{kind} {service}" + ) + + # WQX3 Result defaults to fullPhysChem when no profile is given. + if service == "Result" and not legacy and profile is None: + kwargs["dataProfile"] = "fullPhysChem" + def _is_code_column(name: str) -> bool: """Report whether a WQP column name denotes a code or identifier. @@ -138,8 +184,9 @@ def get_results( US state FIPS code (Example: Illinois is "US:17"). countycode : string US county FIPS code. - huc : string - Eight-digit hydrologic unit (HUC), delimited by semicolons. + huc : string or iterable of strings + Eight-digit hydrologic unit (HUC). Iterable values are encoded as + repeated parameters for WQX3 and semicolon-delimited for legacy WQP. bBox : string Search bounding box (Example: bBox=-92.8,44.2,-88.9,46.0). lat : string @@ -148,8 +195,9 @@ def get_results( Radial-search central longitude in WGS84 decimal degrees. within : string Radial-search distance in decimal miles. - pCode : string - Five-digit USGS parameter code, delimited by semicolons. + pCode : string or iterable of strings + Five-digit USGS parameter code. Iterable values are encoded as + repeated parameters for WQX3 and semicolon-delimited for legacy WQP. NWIS only. startDateLo : string Date of the earliest desired data-collection activity, @@ -157,9 +205,11 @@ def get_results( startDateHi : string Date of the last desired data-collection activity, expressed as 'MM-DD-YYYY'. - characteristicName : string - One or more case-sensitive characteristic names, separated by - semicolons (https://www.waterqualitydata.us/public_srsnames/). + characteristicName : string or iterable of strings + One or more case-sensitive characteristic names + (https://www.waterqualitydata.us/public_srsnames/). Iterable values are + encoded as repeated parameters for WQX3 and semicolon-delimited for + legacy WQP. mimeType : string Output format. Only 'csv' is supported at this time. @@ -199,24 +249,16 @@ def get_results( kwargs = _check_kwargs(kwargs) if legacy is True: - valid_profiles = result_profiles_legacy - kind = "legacy" url = wqp_url("Result") + delimiter = ";" else: - valid_profiles = result_profiles_wqx3 - kind = "WQX3.0" url = wqx3_url("Result") + delimiter = None - profile = kwargs.get("dataProfile") - if profile is not None: - require_one_of( - profile, valid_profiles, name="dataProfile", context=f"{kind} results" - ) - if legacy is not True and profile is None: - kwargs["dataProfile"] = "fullPhysChem" + _validate_profile("Result", legacy, kwargs) response = _query_with_retry( - url, kwargs, delimiter=";", ssl_check=ssl_check, adapter="wqp" + url, kwargs, delimiter=delimiter, ssl_check=ssl_check, adapter="wqp" ) df = _read_wqp_csv(response.text) @@ -240,14 +282,20 @@ def _what( legacy profile. The CSV response is parsed via :func:`_read_wqp_csv`. """ kwargs = _check_kwargs(kwargs) + _validate_profile(service, legacy, kwargs) - if service in services_wqx3: - url = wqp_url(service) if legacy else wqx3_url(service) - else: + if service not in services_wqx3: url = _legacy_only_url(service, legacy=legacy) + delimiter = ";" + elif legacy: + url = wqp_url(service) + delimiter = ";" + else: + url = wqx3_url(service) + delimiter = None response = _query_with_retry( - url, payload=kwargs, delimiter=";", ssl_check=ssl_check, adapter="wqp" + url, payload=kwargs, delimiter=delimiter, ssl_check=ssl_check, adapter="wqp" ) df = _read_wqp_csv(response.text) return df, WQP_Metadata(response, **kwargs) From 5586e340bd5b10daf0b551c1cea21f1b7ec66a5e Mon Sep 17 00:00:00 2001 From: Elise Hinman Date: Wed, 26 Aug 2026 12:25:24 -0500 Subject: [PATCH 2/5] simplify functions further, what and get are basically the same type of function for different services --- dataretrieval/wqp.py | 108 ++++++++++++++++++------------------------- 1 file changed, 46 insertions(+), 62 deletions(-) diff --git a/dataretrieval/wqp.py b/dataretrieval/wqp.py index 06aba0fa..ea93d083 100644 --- a/dataretrieval/wqp.py +++ b/dataretrieval/wqp.py @@ -150,6 +150,43 @@ def _read_wqp_csv(text: str) -> DataFrame: return pd.read_csv(StringIO(text), delimiter=",", low_memory=False, dtype=str_cols) +def _query_wqp( + service: str, + *, + ssl_check: bool, + legacy: bool, + **kwargs: Any, +) -> tuple[DataFrame, WQP_Metadata]: + """Query a WQP service and return the parsed result. + + ``service`` is the WQP service name (e.g. ``"Result"``, ``"Station"``). + Services with a WQX3.0 equivalent (those in :data:`services_wqx3`) use + :func:`wqx3_url` when ``legacy=False`` and :func:`wqp_url` otherwise. + Legacy-only services route through :func:`_legacy_only_url`, which warns + and falls back to the legacy profile. The CSV response is parsed via + :func:`_read_wqp_csv`. + """ + kwargs = _check_kwargs(kwargs) + _validate_profile(service, legacy, kwargs) + + if service not in services_wqx3: + url = _legacy_only_url(service, legacy=legacy) + delimiter = ";" + elif legacy: + url = wqp_url(service) + delimiter = ";" + else: + url = wqx3_url(service) + delimiter = None + + response = _query_with_retry( + url, payload=kwargs, delimiter=delimiter, ssl_check=ssl_check, adapter="wqp" + ) + df = _read_wqp_csv(response.text) + df = _attach_datetime_columns(df) + return df, WQP_Metadata(response, **kwargs) + + def get_results( ssl_check: bool = True, legacy: bool = True, @@ -245,60 +282,7 @@ def get_results( ... ) """ - - kwargs = _check_kwargs(kwargs) - - if legacy is True: - url = wqp_url("Result") - delimiter = ";" - else: - url = wqx3_url("Result") - delimiter = None - - _validate_profile("Result", legacy, kwargs) - - response = _query_with_retry( - url, kwargs, delimiter=delimiter, ssl_check=ssl_check, adapter="wqp" - ) - - df = _read_wqp_csv(response.text) - df = _attach_datetime_columns(df) - return df, WQP_Metadata(response, **kwargs) - - -def _what( - service: str, - *, - ssl_check: bool, - legacy: bool, - **kwargs: Any, -) -> tuple[DataFrame, WQP_Metadata]: - """Shared implementation for the ``what_*`` metadata search functions. - - ``service`` is the WQP service name (e.g. ``"Station"``). Services with a - WQX3.0 equivalent (those in :data:`services_wqx3`) use :func:`wqx3_url` - when ``legacy=False`` and :func:`wqp_url` otherwise. Legacy-only services - route through :func:`_legacy_only_url`, which warns and falls back to the - legacy profile. The CSV response is parsed via :func:`_read_wqp_csv`. - """ - kwargs = _check_kwargs(kwargs) - _validate_profile(service, legacy, kwargs) - - if service not in services_wqx3: - url = _legacy_only_url(service, legacy=legacy) - delimiter = ";" - elif legacy: - url = wqp_url(service) - delimiter = ";" - else: - url = wqx3_url(service) - delimiter = None - - response = _query_with_retry( - url, payload=kwargs, delimiter=delimiter, ssl_check=ssl_check, adapter="wqp" - ) - df = _read_wqp_csv(response.text) - return df, WQP_Metadata(response, **kwargs) + return _query_wqp("Result", ssl_check=ssl_check, legacy=legacy, **kwargs) def what_sites( @@ -345,7 +329,7 @@ def what_sites( """ - return _what("Station", ssl_check=ssl_check, legacy=legacy, **kwargs) + return _query_wqp("Station", ssl_check=ssl_check, legacy=legacy, **kwargs) def what_organizations( @@ -388,7 +372,7 @@ def what_organizations( """ - return _what("Organization", ssl_check=ssl_check, legacy=legacy, **kwargs) + return _query_wqp("Organization", ssl_check=ssl_check, legacy=legacy, **kwargs) def what_projects( @@ -431,7 +415,7 @@ def what_projects( """ - return _what("Project", ssl_check=ssl_check, legacy=legacy, **kwargs) + return _query_wqp("Project", ssl_check=ssl_check, legacy=legacy, **kwargs) def what_activities( @@ -487,7 +471,7 @@ def what_activities( ... ) """ - return _what("Activity", ssl_check=ssl_check, legacy=legacy, **kwargs) + return _query_wqp("Activity", ssl_check=ssl_check, legacy=legacy, **kwargs) def what_detection_limits( @@ -536,7 +520,7 @@ def what_detection_limits( """ - return _what( + return _query_wqp( "ResultDetectionQuantitationLimit", ssl_check=ssl_check, legacy=legacy, @@ -584,7 +568,7 @@ def what_habitat_metrics( """ - return _what("BiologicalMetric", ssl_check=ssl_check, legacy=legacy, **kwargs) + return _query_wqp("BiologicalMetric", ssl_check=ssl_check, legacy=legacy, **kwargs) def what_project_weights( @@ -632,7 +616,7 @@ def what_project_weights( """ - return _what( + return _query_wqp( "ProjectMonitoringLocationWeighting", ssl_check=ssl_check, legacy=legacy, @@ -685,7 +669,7 @@ def what_activity_metrics( """ - return _what("ActivityMetric", ssl_check=ssl_check, legacy=legacy, **kwargs) + return _query_wqp("ActivityMetric", ssl_check=ssl_check, legacy=legacy, **kwargs) def _validate_service(service: str, valid_services: list[str], profile: str) -> None: From 1ecacfad2e063b1db504246bda8b97aeb233d135 Mon Sep 17 00:00:00 2001 From: Elise Hinman Date: Wed, 26 Aug 2026 12:31:02 -0500 Subject: [PATCH 3/5] change iterable to list and remove explanation user doesn't need to know --- dataretrieval/wqp.py | 17 ++++++----------- 1 file changed, 6 insertions(+), 11 deletions(-) diff --git a/dataretrieval/wqp.py b/dataretrieval/wqp.py index ea93d083..54c038ec 100644 --- a/dataretrieval/wqp.py +++ b/dataretrieval/wqp.py @@ -221,9 +221,8 @@ def get_results( US state FIPS code (Example: Illinois is "US:17"). countycode : string US county FIPS code. - huc : string or iterable of strings - Eight-digit hydrologic unit (HUC). Iterable values are encoded as - repeated parameters for WQX3 and semicolon-delimited for legacy WQP. + huc : string or list of strings + Eight-digit hydrologic unit (HUC). bBox : string Search bounding box (Example: bBox=-92.8,44.2,-88.9,46.0). lat : string @@ -232,21 +231,17 @@ def get_results( Radial-search central longitude in WGS84 decimal degrees. within : string Radial-search distance in decimal miles. - pCode : string or iterable of strings - Five-digit USGS parameter code. Iterable values are encoded as - repeated parameters for WQX3 and semicolon-delimited for legacy WQP. - NWIS only. + pCode : string or list of strings + Five-digit USGS parameter code. startDateLo : string Date of the earliest desired data-collection activity, expressed as 'MM-DD-YYYY'. startDateHi : string Date of the last desired data-collection activity, expressed as 'MM-DD-YYYY'. - characteristicName : string or iterable of strings + characteristicName : string or list of strings One or more case-sensitive characteristic names - (https://www.waterqualitydata.us/public_srsnames/). Iterable values are - encoded as repeated parameters for WQX3 and semicolon-delimited for - legacy WQP. + (https://www.waterqualitydata.us/public_srsnames/). mimeType : string Output format. Only 'csv' is supported at this time. From 2bf869b31b69732108665b4f54e942678dd5513b Mon Sep 17 00:00:00 2001 From: thodson-usgs Date: Wed, 26 Aug 2026 14:20:34 -0500 Subject: [PATCH 4/5] Address WQP refactor review --- dataretrieval/wqp.py | 87 ++++++++++++++++++-------------------------- tests/wqp_test.py | 34 +++++++++++++++++ 2 files changed, 69 insertions(+), 52 deletions(-) diff --git a/dataretrieval/wqp.py b/dataretrieval/wqp.py index 54c038ec..134032ec 100644 --- a/dataretrieval/wqp.py +++ b/dataretrieval/wqp.py @@ -61,7 +61,7 @@ result_profiles_wqx3 = ["basicPhysChem", "fullPhysChem", "narrow"] result_profiles_legacy = ["biological", "narrowResult", "resultPhysChem"] -activity_profiles_legacy = ["activityAll"] +activity_profiles_legacy = ["activity", "activityAll"] services_wqx3 = ["Activity", "Result", "Station"] services_legacy = [ "Activity", @@ -75,50 +75,37 @@ "Station", ] -# Valid dataProfile values per service, keyed by (service, legacy: bool). -# None means the service does not accept a dataProfile parameter at all. -_PROFILES: dict[tuple[str, bool], list[str] | None] = { +# Accepted dataProfile values for the endpoints that support the parameter. +# An absent endpoint/profile pair means dataProfile is unsupported there. +_PROFILE_OPTIONS: dict[tuple[str, bool], list[str]] = { ("Result", True): result_profiles_legacy, ("Result", False): result_profiles_wqx3, ("Activity", True): activity_profiles_legacy, - ("Activity", False): None, - ("Station", True): None, - ("Station", False): None, - ("Organization", True): None, - ("Project", True): None, - ("ActivityMetric", True): None, - ("BiologicalMetric", True): None, - ("ProjectMonitoringLocationWeighting", True): None, - ("ResultDetectionQuantitationLimit", True): None, } -def _validate_profile(service: str, legacy: bool, kwargs: dict[str, Any]) -> None: - """Check that ``dataProfile`` is valid for *service*, or absent when unsupported. - - Raises ``ValueError`` if the service does not accept profiles and one was - passed, or if the profile is not in the service's known list. When the - service is Result with ``legacy=False`` and no profile is given, defaults - to ``'fullPhysChem'`` (mutates *kwargs* in place). - """ +def _validate_profile(endpoint: str, legacy: bool, kwargs: dict[str, Any]) -> None: + """Validate ``dataProfile`` for a WQP endpoint and apply its default.""" profile = kwargs.get("dataProfile") - key = (service, legacy) - valid = _PROFILES.get(key) + valid = _PROFILE_OPTIONS.get((endpoint, legacy)) if valid is None and profile is not None: raise ValueError( - f"The {service!r} service does not accept a dataProfile parameter. " + f"The {endpoint!r} endpoint does not accept a dataProfile parameter. " f"Remove dataProfile={profile!r} from your call." ) if valid is not None and profile is not None: kind = "legacy" if legacy else "WQX3.0" require_one_of( - profile, valid, name="dataProfile", context=f"{kind} {service}" + profile, + valid, + name="dataProfile", + context=f"{kind} {endpoint} endpoint", ) # WQX3 Result defaults to fullPhysChem when no profile is given. - if service == "Result" and not legacy and profile is None: + if endpoint == "Result" and not legacy and profile is None: kwargs["dataProfile"] = "fullPhysChem" @@ -151,39 +138,34 @@ def _read_wqp_csv(text: str) -> DataFrame: def _query_wqp( - service: str, + endpoint: str, *, ssl_check: bool, legacy: bool, **kwargs: Any, ) -> tuple[DataFrame, WQP_Metadata]: - """Query a WQP service and return the parsed result. + """Run one WQP getter query against the selected endpoint. - ``service`` is the WQP service name (e.g. ``"Result"``, ``"Station"``). - Services with a WQX3.0 equivalent (those in :data:`services_wqx3`) use + Endpoints with a WQX3.0 equivalent (those in :data:`services_wqx3`) use :func:`wqx3_url` when ``legacy=False`` and :func:`wqp_url` otherwise. - Legacy-only services route through :func:`_legacy_only_url`, which warns + Legacy-only endpoints route through :func:`_legacy_only_url`, which warns and falls back to the legacy profile. The CSV response is parsed via :func:`_read_wqp_csv`. """ kwargs = _check_kwargs(kwargs) - _validate_profile(service, legacy, kwargs) - - if service not in services_wqx3: - url = _legacy_only_url(service, legacy=legacy) - delimiter = ";" - elif legacy: - url = wqp_url(service) - delimiter = ";" + _validate_profile(endpoint, legacy, kwargs) + + if endpoint in services_wqx3: + url = wqp_url(endpoint) if legacy else wqx3_url(endpoint) else: - url = wqx3_url(service) - delimiter = None + url = _legacy_only_url(endpoint, legacy=legacy) response = _query_with_retry( - url, payload=kwargs, delimiter=delimiter, ssl_check=ssl_check, adapter="wqp" + url, payload=kwargs, delimiter=";", ssl_check=ssl_check, adapter="wqp" ) df = _read_wqp_csv(response.text) - df = _attach_datetime_columns(df) + if endpoint == "Result": + df = _attach_datetime_columns(df) return df, WQP_Metadata(response, **kwargs) @@ -221,8 +203,8 @@ def get_results( US state FIPS code (Example: Illinois is "US:17"). countycode : string US county FIPS code. - huc : string or list of strings - Eight-digit hydrologic unit (HUC). + huc : string + Eight-digit hydrologic unit (HUC), delimited by semicolons. bBox : string Search bounding box (Example: bBox=-92.8,44.2,-88.9,46.0). lat : string @@ -231,17 +213,18 @@ def get_results( Radial-search central longitude in WGS84 decimal degrees. within : string Radial-search distance in decimal miles. - pCode : string or list of strings - Five-digit USGS parameter code. + pCode : string + Five-digit USGS parameter code, delimited by semicolons. + NWIS only. startDateLo : string Date of the earliest desired data-collection activity, expressed as 'MM-DD-YYYY'. startDateHi : string Date of the last desired data-collection activity, expressed as 'MM-DD-YYYY'. - characteristicName : string or list of strings - One or more case-sensitive characteristic names - (https://www.waterqualitydata.us/public_srsnames/). + characteristicName : string + One or more case-sensitive characteristic names, separated by + semicolons (https://www.waterqualitydata.us/public_srsnames/). mimeType : string Output format. Only 'csv' is supported at this time. @@ -810,8 +793,8 @@ def _warn_legacy_use() -> None: def _warn_wqx3_unavailable() -> None: - # stacklevel=4: warn -> _warn_wqx3_unavailable -> _legacy_only_url -> _what - # -> what_*, so the warning is attributed to the public ``what_*`` call. + # stacklevel=4: warn -> _warn_wqx3_unavailable -> _legacy_only_url -> + # _query_wqp -> what_*, so the warning is attributed to the public call. warnings.warn( "WQX3.0 profile not available, returning legacy profile.", UserWarning, diff --git a/tests/wqp_test.py b/tests/wqp_test.py index ecc9828f..f582c289 100644 --- a/tests/wqp_test.py +++ b/tests/wqp_test.py @@ -121,6 +121,40 @@ def test_get_results_WQX3(httpx_mock): assert df["Activity_StartDateTime"].notna().all() +@pytest.mark.parametrize("profile", ["activity", "activityAll"]) +def test_what_activities_accepts_documented_legacy_profiles(httpx_mock, profile): + """Accept both Activity profiles published by the WQP OpenAPI document.""" + httpx_mock.add_response(method="GET", text="ActivityIdentifier\nA\n") + + with pytest.warns(DataCurrencyWarning): + df, _ = what_activities(dataProfile=profile) + + assert df["ActivityIdentifier"].tolist() == ["A"] + assert httpx_mock.get_requests()[-1].url.params["dataProfile"] == profile + + +def test_what_activities_preserves_columns_and_row_order(httpx_mock): + """Sharing request plumbing must not change an established return shape.""" + columns = [ + "ActivityIdentifier", + "ActivityStartDate", + "ActivityStartTime/Time", + "ActivityStartTime/TimeZoneCode", + ] + text = ( + ",".join(columns) + + "\nlater,2024-02-01,10:00:00,CST" + + "\nearlier,2024-01-01,10:00:00,CST\n" + ) + httpx_mock.add_response(method="GET", text=text) + + with pytest.warns(DataCurrencyWarning): + df, _ = what_activities() + + assert df.columns.tolist() == columns + assert df["ActivityIdentifier"].tolist() == ["later", "earlier"] + + @pytest.mark.parametrize( ("builder", "service", "expected", "warning"), [ From d7a1dcb839f5a6bd46b6535318814b962fdaa497 Mon Sep 17 00:00:00 2001 From: thodson-usgs Date: Wed, 26 Aug 2026 14:52:52 -0500 Subject: [PATCH 5/5] Apply /simplify cleanups to the WQP funnel Fold the WQX3 Result default into _PROFILE_RULES so the table is the single source of every per-service dataProfile fact, flatten the validation to one branch per fact, and rename endpoint -> service to match the vocabulary the rest of the module and its error messages use. _resolve_profile now returns the payload like its neighbour _check_kwargs rather than writing into the caller's dict behind a check-shaped name. Tests: drive the profile parametrize off activity_profiles_legacy, and replace the hand-built Activity CSV with a fixture comparison inside test_what_query so the "only Result is post-processed" guard is pinned for all eight what_* wrappers instead of one. Co-Authored-By: Claude Opus 5 (1M context) --- dataretrieval/wqp.py | 82 ++++++++++++++++++++++++++------------------ tests/wqp_test.py | 30 ++++------------ 2 files changed, 55 insertions(+), 57 deletions(-) diff --git a/dataretrieval/wqp.py b/dataretrieval/wqp.py index 134032ec..f5cfd072 100644 --- a/dataretrieval/wqp.py +++ b/dataretrieval/wqp.py @@ -13,7 +13,7 @@ import warnings from dataclasses import dataclass from io import StringIO -from typing import TYPE_CHECKING, Any, ClassVar +from typing import TYPE_CHECKING, Any, ClassVar, NamedTuple import pandas as pd @@ -75,38 +75,49 @@ "Station", ] -# Accepted dataProfile values for the endpoints that support the parameter. -# An absent endpoint/profile pair means dataProfile is unsupported there. -_PROFILE_OPTIONS: dict[tuple[str, bool], list[str]] = { - ("Result", True): result_profiles_legacy, - ("Result", False): result_profiles_wqx3, - ("Activity", True): activity_profiles_legacy, + +class _ProfileRule(NamedTuple): + """The ``dataProfile`` contract for one ``(service, legacy)`` pair.""" + + valid: list[str] + default: str | None = None + + +# The dataProfile vocabulary, and any default, per service and profile. +# An absent key means the service does not accept dataProfile at all. +_PROFILE_RULES: dict[tuple[str, bool], _ProfileRule] = { + ("Result", True): _ProfileRule(result_profiles_legacy), + ("Result", False): _ProfileRule(result_profiles_wqx3, default="fullPhysChem"), + ("Activity", True): _ProfileRule(activity_profiles_legacy), } -def _validate_profile(endpoint: str, legacy: bool, kwargs: dict[str, Any]) -> None: - """Validate ``dataProfile`` for a WQP endpoint and apply its default.""" +def _resolve_profile( + service: str, legacy: bool, kwargs: dict[str, Any] +) -> dict[str, Any]: + """Validate ``dataProfile`` for a WQP service and apply the service default.""" + rule = _PROFILE_RULES.get((service, legacy)) profile = kwargs.get("dataProfile") - valid = _PROFILE_OPTIONS.get((endpoint, legacy)) - if valid is None and profile is not None: + if profile is None: + if rule is not None and rule.default is not None: + kwargs["dataProfile"] = rule.default + return kwargs + + if rule is None: raise ValueError( - f"The {endpoint!r} endpoint does not accept a dataProfile parameter. " + f"The {service!r} service does not accept a dataProfile parameter. " f"Remove dataProfile={profile!r} from your call." ) - if valid is not None and profile is not None: - kind = "legacy" if legacy else "WQX3.0" - require_one_of( - profile, - valid, - name="dataProfile", - context=f"{kind} {endpoint} endpoint", - ) - - # WQX3 Result defaults to fullPhysChem when no profile is given. - if endpoint == "Result" and not legacy and profile is None: - kwargs["dataProfile"] = "fullPhysChem" + kind = "legacy" if legacy else "WQX3.0" + require_one_of( + profile, + rule.valid, + name="dataProfile", + context=f"{kind} {service} service", + ) + return kwargs def _is_code_column(name: str) -> bool: @@ -138,33 +149,36 @@ def _read_wqp_csv(text: str) -> DataFrame: def _query_wqp( - endpoint: str, + service: str, *, ssl_check: bool, legacy: bool, **kwargs: Any, ) -> tuple[DataFrame, WQP_Metadata]: - """Run one WQP getter query against the selected endpoint. + """Run one WQP getter query against the selected service. - Endpoints with a WQX3.0 equivalent (those in :data:`services_wqx3`) use + Services with a WQX3.0 equivalent (those in :data:`services_wqx3`) use :func:`wqx3_url` when ``legacy=False`` and :func:`wqp_url` otherwise. - Legacy-only endpoints route through :func:`_legacy_only_url`, which warns - and falls back to the legacy profile. The CSV response is parsed via + Legacy-only services route through :func:`_legacy_only_url`, which warns + and falls back to the legacy profile. ``dataProfile`` is validated against + :data:`_PROFILE_RULES`, and the CSV response is parsed via :func:`_read_wqp_csv`. """ kwargs = _check_kwargs(kwargs) - _validate_profile(endpoint, legacy, kwargs) + kwargs = _resolve_profile(service, legacy, kwargs) - if endpoint in services_wqx3: - url = wqp_url(endpoint) if legacy else wqx3_url(endpoint) + if service in services_wqx3: + url = wqp_url(service) if legacy else wqx3_url(service) else: - url = _legacy_only_url(endpoint, legacy=legacy) + url = _legacy_only_url(service, legacy=legacy) response = _query_with_retry( url, payload=kwargs, delimiter=";", ssl_check=ssl_check, adapter="wqp" ) df = _read_wqp_csv(response.text) - if endpoint == "Result": + # Only get_results documents the appended DateTime columns and the + # activity-start sort, so the other services keep their parsed shape. + if service == "Result": df = _attach_datetime_columns(df) return df, WQP_Metadata(response, **kwargs) diff --git a/tests/wqp_test.py b/tests/wqp_test.py index f582c289..aec36bda 100644 --- a/tests/wqp_test.py +++ b/tests/wqp_test.py @@ -3,6 +3,7 @@ import pytest from pandas import DataFrame +from pandas.testing import assert_frame_equal import dataretrieval import dataretrieval.wqp as wqp @@ -10,6 +11,7 @@ from dataretrieval.wqp import ( WQP_Metadata, _check_kwargs, + _read_wqp_csv, get_results, what_activities, what_activity_metrics, @@ -121,7 +123,7 @@ def test_get_results_WQX3(httpx_mock): assert df["Activity_StartDateTime"].notna().all() -@pytest.mark.parametrize("profile", ["activity", "activityAll"]) +@pytest.mark.parametrize("profile", wqp.activity_profiles_legacy) def test_what_activities_accepts_documented_legacy_profiles(httpx_mock, profile): """Accept both Activity profiles published by the WQP OpenAPI document.""" httpx_mock.add_response(method="GET", text="ActivityIdentifier\nA\n") @@ -133,28 +135,6 @@ def test_what_activities_accepts_documented_legacy_profiles(httpx_mock, profile) assert httpx_mock.get_requests()[-1].url.params["dataProfile"] == profile -def test_what_activities_preserves_columns_and_row_order(httpx_mock): - """Sharing request plumbing must not change an established return shape.""" - columns = [ - "ActivityIdentifier", - "ActivityStartDate", - "ActivityStartTime/Time", - "ActivityStartTime/TimeZoneCode", - ] - text = ( - ",".join(columns) - + "\nlater,2024-02-01,10:00:00,CST" - + "\nearlier,2024-01-01,10:00:00,CST\n" - ) - httpx_mock.add_response(method="GET", text=text) - - with pytest.warns(DataCurrencyWarning): - df, _ = what_activities() - - assert df.columns.tolist() == columns - assert df["ActivityIdentifier"].tolist() == ["later", "earlier"] - - @pytest.mark.parametrize( ("builder", "service", "expected", "warning"), [ @@ -279,6 +259,10 @@ def test_what_query(httpx_mock, func, service, fixture, profile_column): assert type(df) is DataFrame assert not df.empty assert profile_column in df.columns + # Only get_results post-processes: the shared funnel must hand back each + # what_* response exactly as parsed, with no DateTime columns and no sort. + with open(f"tests/data/{fixture}") as text: + assert_frame_equal(df, _read_wqp_csv(text.read())) _assert_wqp_metadata(md, request_url)