Simplify WQP module to share internal functions

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
 
@@ -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",
@@ -76,6 +76,50 @@
 ]
 
 
+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 _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")
+
+    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 {service!r} service does not accept a dataProfile parameter. "
+            f"Remove dataProfile={profile!r} from your call."
+        )
+
+    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:
     """Report whether a WQP column name denotes a code or identifier.
 
@@ -104,6 +148,48 @@ 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]:
+    """Run one WQP getter query against the selected service.
+
+    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. ``dataProfile`` is validated against
+    :data:`_PROFILE_RULES`, and the CSV response is parsed via
+    :func:`_read_wqp_csv`.
+    """
+    kwargs = _check_kwargs(kwargs)
+    kwargs = _resolve_profile(service, legacy, kwargs)
+
+    # WQX3 rejects semicolon-joined array filters, so only the legacy
+    # interface gets a delimiter; delimiter=None repeats the key instead.
+    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)
+    # 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)
+
+
 def get_results(
     ssl_check: bool = True,
     legacy: bool = True,
@@ -194,69 +280,7 @@ 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"
-
-    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)
-
-    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(
@@ -303,7 +327,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(
@@ -346,7 +370,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(
@@ -389,7 +413,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(
@@ -445,7 +469,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(
@@ -494,7 +518,7 @@ def what_detection_limits(
 
     """
 
-    return _what(
+    return _query_wqp(
         "ResultDetectionQuantitationLimit",
         ssl_check=ssl_check,
         legacy=legacy,
@@ -542,7 +566,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(
@@ -590,7 +614,7 @@ def what_project_weights(
 
     """
 
-    return _what(
+    return _query_wqp(
         "ProjectMonitoringLocationWeighting",
         ssl_check=ssl_check,
         legacy=legacy,
@@ -643,7 +667,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:
@@ -791,8 +815,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,

tests/wqp_test.py

@@ -3,13 +3,15 @@
 
 import pytest
 from pandas import DataFrame
+from pandas.testing import assert_frame_equal
 
 import dataretrieval
 import dataretrieval.wqp as wqp
 from dataretrieval.exceptions import DataCurrencyWarning
 from dataretrieval.wqp import (
     WQP_Metadata,
     _check_kwargs,
+    _read_wqp_csv,
     get_results,
     what_activities,
     what_activity_metrics,
@@ -121,6 +123,18 @@ def test_get_results_WQX3(httpx_mock):
     assert df["Activity_StartDateTime"].notna().all()
 
 
+@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")
+
+    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_wqx3_get_results_repeats_list_query_parameters(httpx_mock):
     """WQX3 array filters use repeated keys rather than semicolons."""
     with open("tests/data/wqp3_results.txt") as text:
@@ -319,6 +333,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)