Extract RDB parsing into a shared dataretrieval.rdb module

thodson-usgs · claude · thodson-usgs · commit ec3b21581788 · 2026-05-06T13:40:45.000-05:00
The USGS RDB tab-separated format is used both by NWIS web services
and by the new Water Data STAC catalog's rating-curve assets. Until
now, only nwis.py had a parser, and waterdata/ratings.py reached
across to import the private nwis._read_rdb. That smell is gone:

- New top-level module dataretrieval.rdb with two public helpers:
  - read_rdb(text, dtypes=None): pure parser. Strips the comment
    block, tab-parses the header row, skips the format-spec row,
    and returns a DataFrame. Forwards optional caller-supplied
    dtype hints to pandas.read_csv (unknown column names are
    silently ignored, so dicts are safe to pass on any RDB).
    Raises ValueError on HTML responses.
  - extract_rdb_comment(text): returns the #-prefixed header
    block as a list of lines (the equivalent of R's comment(df)).

- nwis._read_rdb is now a thin wrapper around rdb.read_rdb that
  applies the NWIS-specific dtype hints (site_no, dec_long_va, ...)
  and runs format_response() for datetime indexing. The private
  symbol's contract is unchanged; existing 30 nwis tests pass.
  Dropped the now-unused StringIO import.

- waterdata.ratings switches its import from nwis._read_rdb to
  rdb.read_rdb / extract_rdb_comment, dropping the local
  _extract_rdb_comment helper. The cross-module-private-import
  comment block goes away — the dependency is now a clean
  same-package import.

7 new unit tests pin the shared parser's behavior: basic shape
parsing, format-spec-row skipping, dtype passthrough, all-comments
empty-result, HTML rejection, and the comment extractor.

Net stat: +177 / -72.

Co-Authored-By: Claude Opus 4.7 (1M context) &lt;noreply@anthropic.com&gt;
diff --git a/dataretrieval/nwis.py b/dataretrieval/nwis.py
@@ -7,12 +7,12 @@
 from __future__ import annotations
 
 import warnings
-from io import StringIO
 from json import JSONDecodeError
 
 import pandas as pd
 import requests
 
+from dataretrieval.rdb import read_rdb as _read_rdb_text
 from dataretrieval.utils import BaseMetadata
 
 from .utils import query
@@ -1017,65 +1017,29 @@ def _read_json(json):
     return merged_df
 
 
-def _read_rdb(rdb):
-    """
-    Convert NWIS rdb table into a ``pandas.dataframe``.
+# NWIS-specific column dtype hints; pandas silently ignores unknown
+# names, so passing the dict to read_rdb is safe even on responses
+# whose columns don't include any of these.
+_NWIS_RDB_DTYPES = {
+    "site_no": str,
+    "dec_long_va": float,
+    "dec_lat_va": float,
+    "parm_cd": str,
+    "parameter_cd": str,
+}
 
-    Parameters
-    ----------
-    rdb: string
-        A string representation of an rdb table
 
-    Returns
-    -------
-    df: ``pandas.dataframe``
-        A formatted pandas data frame
+def _read_rdb(rdb):
+    """Parse an NWIS RDB response and apply NWIS-specific post-processing.
 
+    Thin wrapper around :func:`dataretrieval.rdb.read_rdb` that adds the
+    NWIS column-dtype hints and runs :func:`format_response` (datetime
+    index, multi-site MultiIndex, optional GeoDataFrame).
     """
-    if "<html>" in rdb.lower() or "<!doctype html>" in rdb.lower():
-        raise ValueError(
-            "Received HTML response instead of RDB. This often indicates "
-            "that the service has been moved or is currently unavailable."
-        )
-
-    count = 0
-    lines = rdb.splitlines()
-
-    for line in lines:
-        # ignore comment lines
-        if line.startswith("#"):
-            count = count + 1
-
-        else:
-            break
-
-    if count >= len(lines):
-        # All lines are comments — the service returned no data rows (e.g.
-        # "No sites found matching all criteria").  This is a legitimate empty
-        # result, so return an empty DataFrame rather than raising.
-        return pd.DataFrame()
-
-    fields = lines[count].split("\t")
-    fields = [field.replace(",", "").strip() for field in fields if field.strip()]
-    dtypes = {
-        "site_no": str,
-        "dec_long_va": float,
-        "dec_lat_va": float,
-        "parm_cd": str,
-        "parameter_cd": str,
-    }
-
-    df = pd.read_csv(
-        StringIO(rdb),
-        delimiter="\t",
-        skiprows=count + 2,
-        names=fields,
-        na_values="NaN",
-        dtype=dtypes,
-    )
-
-    df = format_response(df)
-    return df
+    df = _read_rdb_text(rdb, dtypes=_NWIS_RDB_DTYPES)
+    if df.empty:
+        return df
+    return format_response(df)
 
 
 def _check_sites_value_types(sites):
diff --git a/dataretrieval/rdb.py b/dataretrieval/rdb.py
@@ -0,0 +1,89 @@
+"""Format-agnostic parser for the USGS RDB tab-separated text format.
+
+RDB (Relational DataBase) is a USGS-internal text format used by NWIS web
+services and by the Water Data STAC catalog's rating-curve assets. Every
+RDB file has the same shape:
+
+- One or more ``#``-prefixed comment lines carrying provenance metadata
+  (data source, retrieval timestamp, station name, parameter codes, etc.).
+- A tab-separated header row naming each column.
+- A second tab-separated row giving column format specs (e.g. ``5s 15s``);
+  it is informational only and skipped during parsing.
+- Tab-separated data rows.
+
+This module exposes the parsing primitives that both ``dataretrieval.nwis``
+and ``dataretrieval.waterdata.ratings`` use. Callers layer their own
+post-processing (NWIS-specific datetime indexing, ratings-specific
+``df.attrs`` provenance, etc.) on top of the raw frame.
+"""
+
+from __future__ import annotations
+
+from io import StringIO
+
+import pandas as pd
+
+
+def read_rdb(rdb: str, dtypes: dict[str, type] | None = None) -> pd.DataFrame:
+    """Parse an RDB text response into a ``pandas.DataFrame``.
+
+    Parameters
+    ----------
+    rdb : str
+        The RDB text response from a USGS web service.
+    dtypes : dict[str, type] or None, optional
+        Optional column-name to dtype hints, forwarded to
+        ``pandas.read_csv``. Unknown column names are silently ignored, so
+        callers may safely pass a dict of all columns they might be
+        interested in.
+
+    Returns
+    -------
+    pandas.DataFrame
+        The parsed data. An RDB consisting only of comment lines (e.g. a
+        "no sites found" response) returns an empty DataFrame rather than
+        raising.
+
+    Raises
+    ------
+    ValueError
+        If the response body looks like HTML, which usually means the
+        service has been moved, is degraded, or returned an error page.
+    """
+    if "<html>" in rdb.lower() or "<!doctype html>" in rdb.lower():
+        raise ValueError(
+            "Received HTML response instead of RDB. This often indicates "
+            "that the service has been moved or is currently unavailable."
+        )
+
+    lines = rdb.splitlines()
+    header_idx = next(
+        (i for i, line in enumerate(lines) if not line.startswith("#")),
+        len(lines),
+    )
+    if header_idx >= len(lines):
+        # All lines are comments — a legitimate empty result.
+        return pd.DataFrame()
+
+    fields = lines[header_idx].split("\t")
+    fields = [f.replace(",", "").strip() for f in fields if f.strip()]
+
+    return pd.read_csv(
+        StringIO(rdb),
+        delimiter="\t",
+        skiprows=header_idx + 2,  # +1 for header, +1 for the format-spec row
+        names=fields,
+        na_values="NaN",
+        dtype=dtypes,
+    )
+
+
+def extract_rdb_comment(rdb: str) -> list[str]:
+    """Return the RDB ``#``-prefixed comment block as a list of header lines.
+
+    The comment block carries provenance metadata that is otherwise lost
+    during parsing — data source, retrieval timestamp, parameter codes,
+    rating id and last-shifted timestamp for ratings, etc. R's
+    ``dataRetrieval`` exposes the equivalent via ``comment(df)``.
+    """
+    return [line for line in rdb.splitlines() if line.startswith("#")]
diff --git a/dataretrieval/waterdata/ratings.py b/dataretrieval/waterdata/ratings.py
@@ -19,10 +19,7 @@
 import pandas as pd
 import requests
 
-# Rating files use the same USGS RDB shape as NWIS responses, so we reuse
-# the parser already in ``nwis``. ``_read_rdb`` is private; if it ever
-# moves we want a loud failure here, hence the explicit import.
-from dataretrieval.nwis import _read_rdb
+from dataretrieval.rdb import extract_rdb_comment, read_rdb
 
 from .utils import _DURATION_RE, BASE_URL, _default_headers, _format_api_dates
 
@@ -245,16 +242,6 @@ def _search(
     return response.json().get("features", [])
 
 
-def _extract_rdb_comment(rdb: str) -> list[str]:
-    """Return the RDB ``#``-prefixed comment block as a list of header lines.
-
-    Carries useful per-rating provenance — rating id, parameter description,
-    expansion type, last-shifted timestamp, etc. R's ``read_waterdata_ratings``
-    exposes the equivalent via ``comment(df)``.
-    """
-    return [line for line in rdb.splitlines() if line.startswith("#")]
-
-
 def _download_and_parse(
     feature: dict[str, Any],
     file_path: str | None,
@@ -269,7 +256,7 @@ def _download_and_parse(
         with open(os.path.join(file_path, feature["id"]), "w") as f:
             f.write(response.text)
 
-    df = _read_rdb(response.text)
-    df.attrs["comment"] = _extract_rdb_comment(response.text)
+    df = read_rdb(response.text)
+    df.attrs["comment"] = extract_rdb_comment(response.text)
     df.attrs["url"] = url
     return df
diff --git a/tests/rdb_test.py b/tests/rdb_test.py
@@ -0,0 +1,65 @@
+import pandas as pd
+import pytest
+
+from dataretrieval.rdb import extract_rdb_comment, read_rdb
+
+# A minimally complete RDB: comment block, header row, format-spec row,
+# data rows. Both NWIS responses and ratings RDBs share this shape.
+_BASIC_RDB = """\
+# header line one
+# header line two
+agency_cd\tsite_no\tINDEP\tDEP
+5s\t15s\t10n\t10n
+USGS\t01104475\t0.10\t0.0
+USGS\t01104475\t0.20\t0.5
+USGS\t01104475\t0.30\t1.2
+"""
+
+
+def test_read_rdb_parses_basic_shape():
+    df = read_rdb(_BASIC_RDB)
+    assert list(df.columns) == ["agency_cd", "site_no", "INDEP", "DEP"]
+    assert len(df) == 3
+    assert df["INDEP"].tolist() == [0.10, 0.20, 0.30]
+
+
+def test_read_rdb_skips_format_spec_row():
+    """The "5s 15s 10n 10n" row is metadata, not data."""
+    df = read_rdb(_BASIC_RDB)
+    # If the format-spec row had been treated as data, df would have 4 rows
+    # and "5s" / "15s" would appear in the parsed values.
+    assert "5s" not in df["agency_cd"].tolist()
+
+
+def test_read_rdb_dtype_hints_applied():
+    """Caller-supplied dtype hints are forwarded to pandas; unknown names ignored."""
+    df = read_rdb(_BASIC_RDB, dtypes={"site_no": str, "DEP": float, "unknown": int})
+    # Without the str hint, pandas would parse "01104475" as int and drop the
+    # leading zero. Check the values, not the dtype name (which varies across
+    # pandas versions: object, StringDtype, etc.).
+    assert df["site_no"].iloc[0] == "01104475"
+    assert df["DEP"].dtype == float
+
+
+def test_read_rdb_empty_when_only_comments():
+    """All-comments input is a legitimate "no data" response, not an error."""
+    df = read_rdb("# only a comment\n# and another\n")
+    assert isinstance(df, pd.DataFrame)
+    assert df.empty
+
+
+def test_read_rdb_raises_on_html_response():
+    """If the service returns an HTML error page, surface it loudly."""
+    with pytest.raises(ValueError, match="HTML"):
+        read_rdb("<html><body>Service Unavailable</body></html>")
+    with pytest.raises(ValueError, match="HTML"):
+        read_rdb("<!DOCTYPE html>\n<html>...")
+
+
+def test_extract_rdb_comment_returns_only_hash_lines():
+    comments = extract_rdb_comment(_BASIC_RDB)
+    assert comments == ["# header line one", "# header line two"]
+
+
+def test_extract_rdb_comment_empty_when_no_comments():
+    assert extract_rdb_comment("a\tb\n1\t2\n") == []
