diff --git a/dataretrieval/_configuration_core.py b/dataretrieval/_configuration_core.py index d1b1b460..bfe40272 100644 --- a/dataretrieval/_configuration_core.py +++ b/dataretrieval/_configuration_core.py @@ -1113,16 +1113,9 @@ def _adapter_file_settings( def _load_file(path: Path) -> _ParsedFile: """Parse the configuration file at *path*, caching until it changes on disk.""" global _file_cache - try: - st = path.stat() - except FileNotFoundError: - # No file is the normal case: continue to the built-in default. One - # shared empty instance rather than a fresh one per read -- nothing - # mutates a ``_ParsedFile``, and returning the same object each time is - # what lets callers memoize on its identity. + st = _stat_config_file(path) + if st is None: return _NO_FILE - except OSError as exc: - raise ConfigurationError(f"could not access {path}: {exc}") from exc if stat.S_ISDIR(st.st_mode): raise ConfigurationError( @@ -1134,32 +1127,42 @@ def _load_file(path: Path) -> _ParsedFile: # opened, which is what ``DATARETRIEVAL_CONFIG=/dev/null`` asks for and the # only coherent answer for a stream: settings are re-resolved on every # request, so a FIFO would hand its contents to the first getter and - # nothing to the rest, making the API key vanish mid-run. (It would also - # block on open until a writer appeared.) + # nothing to the rest, making the API key vanish mid-run. if not stat.S_ISREG(st.st_mode): return _ParsedFile(exists=True) - # POSIX ``st_ctime_ns`` advances on any inode change, so the metadata stamp - # catches even a rewrite that restores the original mtime (``cp -p``, rsync - # ``--times``, an editor that preserves timestamps). Windows ctime is - # *creation* time, so there the stamp cannot see that class of edit and the - # content compare below is the only correct check -- worth the re-read, - # since serving a stale API key is the alternative. - # - # Dropping this gate (or dropping ctime from the stamp so Windows can use - # it) has been proposed repeatedly on the grounds that the re-read is - # wasteful. It is, but it is also the only thing standing between a - # timestamp-preserving write and a stale credential; a ctime-less stamp is - # identical across exactly that edit. ``test_file_edit_is_picked_up`` - # pins the behavior. Please do not "optimize" it without a Windows-safe - # change detector. - # - # Measured, so the next reviewer does not have to re-derive it: forcing the - # Windows branch costs 27 us per settings read against 5 us with the stamp - # (5 syscalls instead of 1; a 64-byte file and a 6.5 kB one measure the - # same, since the content compare still spares the TOML parse). At the 8 - # reads a one-chunk query performs that is ~175 us against a 100-500 ms - # round trip -- 0.04%, and the alternative is serving a stale key. + cached_parse = _cached_parse_by_metadata(path, st) + if cached_parse is not None: + return cached_parse + + content, opened_st = _read_file_content(path) + parsed = _parse_or_reuse_cache(path, content) + _warn_on_loose_permissions(path, opened_st, parsed) + _file_cache = (path, _file_stamp(opened_st), content, parsed) + return parsed + + +def _stat_config_file(path: Path) -> os.stat_result | None: + """Stat the config file, returning ``None`` if it does not exist.""" + try: + return path.stat() + except FileNotFoundError: + return None + except OSError as exc: + raise ConfigurationError(f"could not access {path}: {exc}") from exc + + +def _cached_parse_by_metadata(path: Path, st: os.stat_result) -> _ParsedFile | None: + """The cached parse when the metadata stamp still matches, else ``None``. + + POSIX ``st_ctime_ns`` advances on any inode change, so the metadata stamp + catches even a rewrite that restores the original mtime. Windows ctime is + *creation* time, so there the stamp cannot see that class of edit and the + content compare in :func:`_parse_or_reuse_cache` is the only correct check + -- the re-read it forces is deliberate, and ``test_file_edit_is_picked_up`` + pins it. Do not drop the ctime gate (or extend the stamp to Windows) + without a Windows-safe change detector. + """ cached = _file_cache if ( os.name != "nt" @@ -1168,28 +1171,33 @@ def _load_file(path: Path) -> _ParsedFile: and cached[1] == _file_stamp(st) ): return cached[3] + return None + +def _read_file_content(path: Path) -> tuple[bytes, os.stat_result]: + """Read the file content and return it with the stat of the opened handle.""" try: with path.open("rb") as handle: content = handle.read() opened_st = os.fstat(handle.fileno()) except OSError as exc: raise ConfigurationError(f"could not read {path}: {exc}") from exc + return content, opened_st + +def _parse_or_reuse_cache(path: Path, content: bytes) -> _ParsedFile: + """Parse the TOML content, reusing the cache if content is unchanged.""" + cached = _file_cache if cached is not None and cached[0] is path and cached[2] == content: - parsed = cached[3] - else: - tomllib = _toml_parser() - try: - data = tomllib.loads(content.decode("utf-8")) - except UnicodeDecodeError as exc: - raise ConfigurationError(f"{path} is not valid UTF-8: {exc}") from exc - except tomllib.TOMLDecodeError as exc: - raise ConfigurationError(f"{path} is not valid TOML: {exc}") from exc - parsed = _interpret(data, path) - _warn_on_loose_permissions(path, opened_st, parsed) - _file_cache = (path, _file_stamp(opened_st), content, parsed) - return parsed + return cached[3] + tomllib = _toml_parser() + try: + data = tomllib.loads(content.decode("utf-8")) + except UnicodeDecodeError as exc: + raise ConfigurationError(f"{path} is not valid UTF-8: {exc}") from exc + except tomllib.TOMLDecodeError as exc: + raise ConfigurationError(f"{path} is not valid TOML: {exc}") from exc + return _interpret(data, path) def _file_stamp(st: os.stat_result) -> _FileStamp: diff --git a/dataretrieval/_wqx.py b/dataretrieval/_wqx.py index 8089c68d..dc21afa7 100644 --- a/dataretrieval/_wqx.py +++ b/dataretrieval/_wqx.py @@ -48,6 +48,49 @@ def _build_utc_datetime( ) +def _build_triplet_datetime( + df: pd.DataFrame, prefix: str, columns: set[str] +) -> pd.Series | None: + """Try each Time/TimeZone suffix pair and return a UTC Series, or None.""" + for time_suffix, tz_suffix in _TIME_TZ_SUFFIXES: + time_col = prefix + time_suffix + tz_col = prefix + tz_suffix + if time_col in columns and tz_col in columns: + return _build_utc_datetime(df[prefix + "Date"], df[time_col], df[tz_col]) + return None + + +def _find_datetime_triplets(df: pd.DataFrame) -> dict[str, pd.Series]: + """Detect Date/Time/TimeZone column triplets and build UTC datetime columns. + + Returns a mapping of new column names to UTC Series. + """ + columns = set(df.columns) + new_columns: dict[str, pd.Series] = {} + + for col in df.columns: + if not col.endswith("Date"): + continue + prefix = col.removesuffix("Date") + target = prefix + "DateTime" + if target in columns or target in new_columns: + continue + utc_series = _build_triplet_datetime(df, prefix, columns) + if utc_series is not None: + new_columns[target] = utc_series + + return new_columns + + +def _resolve_sort_key(df: pd.DataFrame) -> str | None: + """Pick the canonical sort column, falling back to the first ``*Date``.""" + if "Activity_StartDateTime" in df.columns: + return "Activity_StartDateTime" + if "ActivityStartDateTime" in df.columns: + return "ActivityStartDateTime" + return next((c for c in df.columns if c.endswith("Date")), None) + + def _attach_datetime_columns(df: pd.DataFrame) -> pd.DataFrame: """Append a UTC ``DateTime`` column per Date/Time/TimeZone triplet. @@ -82,37 +125,16 @@ def _attach_datetime_columns(df: pd.DataFrame) -> pd.DataFrame: and rows sorted by the activity-start datetime (if any date column was detected). """ - columns = set(df.columns) - new_columns = {} - first_date_col = None - for col in df.columns: - if not col.endswith("Date"): - continue - if first_date_col is None: - first_date_col = col - prefix = col.removesuffix("Date") - target = prefix + "DateTime" - if target in columns or target in new_columns: - continue - for time_suffix, tz_suffix in _TIME_TZ_SUFFIXES: - time_col = prefix + time_suffix - tz_col = prefix + tz_suffix - if time_col in columns and tz_col in columns: - new_columns[target] = _build_utc_datetime( - df[col], df[time_col], df[tz_col] - ) - break + new_columns = _find_datetime_triplets(df) + if new_columns: # Concat in one shot — per-column assignment on a wide CSV-derived # frame triggers pandas' fragmentation PerformanceWarning. df = pd.concat([df, pd.DataFrame(new_columns, index=df.index)], axis=1) - sort_key: str | None - if "Activity_StartDateTime" in df.columns: - sort_key = "Activity_StartDateTime" - elif "ActivityStartDateTime" in df.columns: - sort_key = "ActivityStartDateTime" - else: - sort_key = first_date_col + + # The appended columns end in "DateTime", so the first "*Date" column is + # the same before and after the concat. + sort_key = _resolve_sort_key(df) if sort_key is not None: df = df.sort_values(by=sort_key, ignore_index=True) return df diff --git a/dataretrieval/codes/states.py b/dataretrieval/codes/states.py index 79e0a18d..321781b2 100644 --- a/dataretrieval/codes/states.py +++ b/dataretrieval/codes/states.py @@ -180,6 +180,11 @@ def _to_state_one(value: str | int, to: str) -> str: f'code ("WI"), or a two-digit ANSI/FIPS code ("55").' ) + return _format_state(name, to) + + +def _format_state(name: str, to: str) -> str: + """Render a canonical state *name* in the ``to`` representation.""" if to == "name": return name if to == "postal": diff --git a/dataretrieval/configuration.py b/dataretrieval/configuration.py index 93a18455..4c6992bc 100644 --- a/dataretrieval/configuration.py +++ b/dataretrieval/configuration.py @@ -76,6 +76,7 @@ from collections.abc import Callable, Iterator, Mapping from contextlib import contextmanager from functools import partial +from pathlib import Path from typing import TextIO, overload # Explicit same-name aliases preserve the facade's public and private compatibility @@ -246,6 +247,7 @@ def _frame(configurations: tuple[BaseConfiguration, ...]) -> _Frame: "Configuration(...); a setting for one service goes on that " "adapter's configuration, e.g. WaterdataConfiguration(...)." ) + adapter = configuration.adapter if adapter in seen: where = f"the {adapter} adapter" if adapter else "the package-wide settings" @@ -253,16 +255,27 @@ def _frame(configurations: tuple[BaseConfiguration, ...]) -> _Frame: f"configure() got two configurations for {where}. Precedence " "between them would be undefined, so combine them into one." ) + seen.add(adapter) - label = configuration._provenance() - for name, value in configuration.values().items(): - key: _ScopeKey = name if adapter is None else (adapter, name) - raw = ( - None - if value is None - else _coerce_typed(name, value, configuration._source(name)) - ) - overrides[key] = (raw, label) + overrides.update(_configuration_overrides(configuration)) + return overrides + + +def _configuration_overrides( + configuration: BaseConfiguration, +) -> dict[_ScopeKey, tuple[_SettingValue, str]]: + """Render one configuration's values into raw-string overrides.""" + adapter = configuration.adapter + label = configuration._provenance() + overrides: dict[_ScopeKey, tuple[_SettingValue, str]] = {} + for name, value in configuration.values().items(): + key: _ScopeKey = name if adapter is None else (adapter, name) + raw = ( + None + if value is None + else _coerce_typed(name, value, configuration._source(name)) + ) + overrides[key] = (raw, label) return overrides @@ -321,54 +334,79 @@ def show_configuration(*, stream: TextIO | None = None) -> None: print(f"config file ", file=out) return - # Nothing here raises. This function exists to explain a configuration, and - # the configurations most in need of explaining are the broken ones -- an - # unparseable file, a value that fails its grammar, a profile that no - # longer exists. Each distinct failure is printed once, in the first place - # it shows up; a repeat is collapsed, so one bad file does not bury the - # rows that did resolve under ten copies of the same message. - reported: str | None = None + cell = _ErrorDeduplicatingCell() + parsed = _show_file_status(out, path, cell) - def cell(render: Callable[[], object]) -> str: - nonlocal reported + rows = [ + (name, cell(partial(_DISPLAYS[name], None)), cell(partial(_source_label, name))) + for name in SETTINGS + ] + _print_setting_rows(out, rows) + _show_built_in_default_note(out, rows) + _show_adapter_overrides(out, cell, {name: source for name, _value, source in rows}) + _show_profiles(out, parsed) + _show_unimported_adapters(out) + + +class _ErrorDeduplicatingCell: + """Render a value, deduplicating consecutive configuration errors. + + The report exists to explain a configuration, and the configurations + most in need of explaining are the broken ones -- an unparseable file, a + value that fails its grammar, a profile that no longer exists. Nothing + here raises: each distinct failure is printed once, in the first place + it shows up; a repeat is collapsed, so one bad file does not bury the + rows that did resolve under ten copies of the same message. + """ + + def __init__(self) -> None: + self._reported: str | None = None + + def mark_reported(self, exc: ConfigurationError) -> None: + """Record that *exc* was already printed, so a repeat collapses.""" + self._reported = str(exc) + + def __call__(self, render: Callable[[], object]) -> str: try: value = render() except ConfigurationError as exc: - if str(exc) == reported: + if str(exc) == self._reported: return "" - reported = str(exc) + self._reported = str(exc) return f"" return "" if value is None else str(value) - # Probing the file once here means a whole-file problem -- unparseable - # TOML, a bad value at the top level -- is reported on the file row rather - # than repeated in every setting's row below. The parsed form is kept for - # the profile section, which asks what the file *defines* rather than what - # resolved; an unparseable file defines nothing, and has already said so - # here. + +def _show_file_status( + out: TextIO, path: Path, cell: _ErrorDeduplicatingCell +) -> _ParsedFile: + """Probe and print the config file status line, returning the parsed file. + + Probing the file once here means a whole-file problem -- unparseable TOML, + a bad value at the top level -- is reported on the file row rather than + repeated in every setting's row below. + """ parsed = _NO_FILE try: _, parsed = _current_file() status = "found" if path.exists() else "not found" except ConfigurationError as exc: - reported = str(exc) + cell.mark_reported(exc) status = f"ERROR: {exc}" print(f"config file {path} ({status})", file=out) + return parsed - rows = [ - (name, cell(partial(_DISPLAYS[name], None)), cell(partial(_source_label, name))) - for name in SETTINGS - ] + +def _print_setting_rows(out: TextIO, rows: list[tuple[str, str, str]]) -> None: + """Print the package-wide setting rows in aligned columns.""" name_width = max(len(name) for name, _value, _source in rows) value_width = max(len(value) for _name, value, _source in rows) for name, value, source in rows: print(f"{name:<{name_width}} {value:<{value_width}} {source}", file=out) - # A built-in default is package-wide, and a service may prefer its own for - # its own calls -- so a row reading "built-in default" is not a promise - # about every service. Saying so is the honest scope of this report: this - # module is a leaf and cannot enumerate the services, and a value from any - # source outranks both kinds of default anyway. + +def _show_built_in_default_note(out: TextIO, rows: list[tuple[str, str, str]]) -> None: + """Print the built-in default footnote when at least one row uses it.""" if any(source == _BUILT_IN for _name, _value, source in rows): print( "\nA built-in default is package-wide. An adapter may prefer its own " @@ -376,10 +414,6 @@ def cell(render: Callable[[], object]) -> str: file=out, ) - _show_adapter_overrides(out, cell, {name: source for name, _value, source in rows}) - _show_profiles(out, parsed) - _show_unimported_adapters(out) - def _show_adapter_overrides( out: TextIO, @@ -400,37 +434,60 @@ def _show_adapter_overrides( against, so it is skipped here and named by :func:`_show_unimported_adapters` instead. """ + overrides = _collect_adapter_overrides(cell, package_wide) + if not overrides: + return + print("\nadapter overrides", file=out) + a_width = max(len(a) for a, _n, _v, _s in overrides) + n_width = max(len(n) for _a, n, _v, _s in overrides) + v_width = max(len(v) for _a, _n, v, _s in overrides) + for adapter, name, value, source in overrides: + print( + f" {adapter:<{a_width}} {name:<{n_width}} {value:<{v_width}} {source}", + file=out, + ) + + +def _collect_adapter_overrides( + cell: Callable[[Callable[[], object]], str], + package_wide: Mapping[str, str], +) -> list[tuple[str, str, str, str]]: + """Gather adapter-scoped settings that differ from the package-wide rows. + + Separated from :func:`_show_adapter_overrides` so the collection logic -- + which carries the nesting -- is not interleaved with the formatting logic. + """ overrides: list[tuple[str, str, str, str]] = [] for adapter in ADAPTERS: accepted = settings_for(adapter) if accepted is None: continue - for name in _ALL_SETTINGS: - if name not in accepted: - continue - scoped = cell(partial(_source_label, name, adapter)) - # ``package_wide`` is what the rows above already resolved. Asking - # again would repeat the work once per adapter *and* consume the - # shared error-dedupe state, so a broken config's message could be - # collapsed here before the row that needs it prints. An - # adapter-only setting has no row above, and no package-wide value - # it could inherit, so its baseline is the built-in default. - if scoped == package_wide.get(name, _BUILT_IN): - continue # inherited from the package-wide tier - value = cell(partial(_DISPLAYS[name], adapter)) - overrides.append((adapter, name, value, scoped)) - - if overrides: - print("\nadapter overrides", file=out) - a_width = max(len(a) for a, _n, _v, _s in overrides) - n_width = max(len(n) for _a, n, _v, _s in overrides) - v_width = max(len(v) for _a, _n, v, _s in overrides) - for adapter, name, value, source in overrides: - print( - f" {adapter:<{a_width}} {name:<{n_width}} " - f"{value:<{v_width}} {source}", - file=out, - ) + overrides.extend(_overrides_for_adapter(adapter, accepted, cell, package_wide)) + return overrides + + +def _overrides_for_adapter( + adapter: str, + accepted: frozenset[str], + cell: Callable[[Callable[[], object]], str], + package_wide: Mapping[str, str], +) -> list[tuple[str, str, str, str]]: + """The overridden settings for one adapter. + + ``package_wide`` is what the rows above already resolved. An adapter-only + setting has no row above, and no package-wide value it could inherit, so + its baseline is the built-in default. + """ + overrides: list[tuple[str, str, str, str]] = [] + for name in _ALL_SETTINGS: + if name not in accepted: + continue + scoped = cell(partial(_source_label, name, adapter)) + if scoped == package_wide.get(name, _BUILT_IN): + continue # inherited from the package-wide tier + value = cell(partial(_DISPLAYS[name], adapter)) + overrides.append((adapter, name, value, scoped)) + return overrides def _show_profiles(out: TextIO, parsed: _ParsedFile) -> None: @@ -639,25 +696,53 @@ def _resolve(name: str, adapter: str | None = None) -> tuple[str | None, str, st tier answered (one of the constants above) -- ``None`` with ``_BUILT_IN`` / ``_DEFAULT`` when nothing configured it. """ - # An adapter name nobody recognizes is a typo in *our* source, and its - # failure mode is silence: ``_accepts`` would wave every setting through, - # the file would hold no table under that name, and the read would fall - # through to the package-wide value -- so a ``[waterdata]`` table, or a - # ``WaterdataConfiguration``, would be ignored with nothing raised - # anywhere. Checked here rather than left to the fitness test that greps - # for ``adapter=""``, which can only see that the string occurs. - if adapter is not None and adapter not in ADAPTERS: + _check_adapter_known(adapter) + _check_env_not_refused(name) + + # ``None`` unless this adapter actually reads this setting, so a setting + # outside its vocabulary resolves package-wide rather than looking for a + # scope it could never have been written into. + scoped: str | None = ( + adapter if adapter is not None and _accepts(adapter, name) else None + ) + + from_block = _resolve_from_block(name, scoped) + if from_block is not None: + return from_block + + from_env = _resolve_from_env(name) + if from_env is not None: + return from_env + + return _resolve_from_file(name, scoped) + + +def _check_adapter_known(adapter: str | None) -> None: + """Raise if *adapter* is not in the configurable adapter roster. + + An adapter name nobody recognizes is a typo in *our* source, and its + failure mode is silence: ``_accepts`` would wave every setting through, + the file would hold no table under that name, and the read would fall + through to the package-wide value -- so a ``[waterdata]`` table, or a + ``WaterdataConfiguration``, would be ignored with nothing raised anywhere. + """ + if adapter not in (*ADAPTERS, None): raise ConfigurationError( f"{adapter!r} is not a configurable adapter. The adapters are " f"{', '.join(ADAPTERS)}." ) - # Refused before anything is consulted, not at the environment's turn in - # the chain. The file refuses ``base_url`` whether or not a block also set - # one -- it raises while the file is read -- and the two surfaces are one - # rule, so a variable that cannot work must not be silently outranked by a - # block that happens to work. Unsetting it is the only fix, and the message - # says so. + +def _check_env_not_refused(name: str) -> None: + """Raise if an environment variable is set for a code-only setting. + + Refused before anything is consulted, not at the environment's turn in + the chain. The file refuses ``base_url`` whether or not a block also set + one -- it raises while the file is read -- and the two surfaces are one + rule, so a variable that cannot work must not be silently outranked by a + block that happens to work. Unsetting it is the only fix, and the message + says so. + """ refused = _REFUSED_ENV_VARS.get(name) if refused is not None and refused in os.environ: raise ConfigurationError( @@ -667,37 +752,48 @@ def _resolve(name: str, adapter: str | None = None) -> tuple[str | None, str, st f"WaterdataConfiguration({name}=...)." ) - # ``None`` unless this adapter actually reads this setting, so a setting - # outside its vocabulary resolves package-wide rather than looking for a - # scope it could never have been written into. - scoped: str | None = ( - adapter if adapter is not None and _accepts(adapter, name) else None - ) - # Innermost block first: a value set by a nested block wins over both - # scopes of an enclosing one. Within one block the adapter-scoped value is - # the more specific of the two, so it is asked first. Each entry already - # carries its own label, which is what keeps the profile a value came from - # reportable (:data:`_Frame`). +def _resolve_from_block( + name: str, scoped: str | None +) -> tuple[str | None, str, str] | None: + """Walk the scope stack for the first block that sets *name*. + + Innermost block first: a value set by a nested block wins over both + scopes of an enclosing one. Within one block the adapter-scoped value is + the more specific of the two, so it is asked first. + """ for frame in reversed(_scope.get()): if scoped is not None and (scoped, name) in frame: return (*frame[(scoped, name)], _BLOCK) if name in frame: return (*frame[name], _BLOCK) + return None - # No per-adapter environment variables: seven adapters times four settings - # is a namespace nobody can hold in mind, and an exported variable is - # invisible at the call site. See ADR 0010. + +def _resolve_from_env(name: str) -> tuple[str | None, str, str] | None: + """Check whether an environment variable supplies the setting. + + No per-adapter environment variables: seven adapters times four settings + is a namespace nobody can hold in mind, and an exported variable is + invisible at the call site. See ADR 0010. + """ env = ENV_VARS.get(name) - if env is not None: - raw = os.environ.get(env) - if raw is not None and (raw.strip() or name in _BLANK_MEANS_SET): - return raw, _env_source_label(env), _ENV - - # One load serves both file tiers. Reading the file twice -- once for the - # adapter table, once for the top level -- cost a second stat on every - # adapter-scoped resolution, and the common case (no table for this - # adapter) is the one that paid it. + if env is None: + return None + raw = os.environ.get(env) + if raw is not None and (raw.strip() or name in _BLANK_MEANS_SET): + return raw, _env_source_label(env), _ENV + return None + + +def _resolve_from_file(name: str, scoped: str | None) -> tuple[str | None, str, str]: + """Fall through to the configuration file, then the built-in default. + + One load serves both file tiers. Reading the file twice -- once for the + adapter table, once for the top level -- cost a second stat on every + adapter-scoped resolution, and the common case (no table for this adapter) + is the one that paid it. + """ path, parsed = _current_file() if scoped is not None: diff --git a/dataretrieval/interruptions.py b/dataretrieval/interruptions.py index 0bd0d0cc..daa1510f 100644 --- a/dataretrieval/interruptions.py +++ b/dataretrieval/interruptions.py @@ -130,26 +130,13 @@ def __init__( retry_after: float | None = None, cause: BaseException | None = None, ) -> None: - message = self._MESSAGE_TEMPLATE.format( - completed_chunks=completed_chunks, total_chunks=total_chunks - ) - if cause is not None: - cause_msg = str(cause) or type(cause).__name__ - message = f"{message} Cause: {type(cause).__name__}: {cause_msg}" + message = self._format_message(completed_chunks, total_chunks, cause) super().__init__(message) self.completed_chunks = completed_chunks self.total_chunks = total_chunks self.call = call self.retry_after = retry_after - self.status_code = getattr(type(self), "_DEFAULT_STATUS", None) - if self.status_code is None and cause is not None: - # The status is usually a few frames down: a typed error raised - # ``from`` the httpx failure that carried it. - for current in _walk_causes(cause): - status = getattr(current, "status_code", None) - if status is not None: - self.status_code = status - break + self.status_code = self._resolve_status_code(cause) # Snapshot partial state at raise time so the exception stays a stable # record of the failure moment: ``exc.partial_frame`` / # ``.partial_response`` do NOT advance on a later ``call.resume()`` @@ -165,6 +152,34 @@ def __init__( self.partial_frame = call.partial_frame.copy() self.partial_response = call.partial_response + def _format_message( + self, + completed_chunks: int, + total_chunks: int, + cause: BaseException | None, + ) -> str: + """Build the exception message from the template, appending cause info.""" + message = self._MESSAGE_TEMPLATE.format( + completed_chunks=completed_chunks, total_chunks=total_chunks + ) + if cause is not None: + cause_msg = str(cause) or type(cause).__name__ + message = f"{message} Cause: {type(cause).__name__}: {cause_msg}" + return message + + def _resolve_status_code(self, cause: BaseException | None) -> int | None: + """Resolve the HTTP status code from the class default or cause chain.""" + status: int | None = getattr(type(self), "_DEFAULT_STATUS", None) + if status is not None or cause is None: + return status + # The status is usually a few frames down: a typed error raised + # ``from`` the httpx failure that carried it. + for current in _walk_causes(cause): + found: int | None = getattr(current, "status_code", None) + if found is not None: + return found + return None + def __getstate__(self) -> dict[str, Any]: # Drop the live FanOut before pickling: its ``.fetch`` is an # undecorated module function pickle can't reference by name, so the diff --git a/dataretrieval/nldi.py b/dataretrieval/nldi.py index bbae1c11..1649b305 100644 --- a/dataretrieval/nldi.py +++ b/dataretrieval/nldi.py @@ -326,6 +326,23 @@ def _navigation_request( return url, {"distance": str(distance)} +def _validate_lat_long_origin( + comid: int | None, + feature_source: str | None, + feature_id: str | None, +) -> None: + """Raise if lat/long is combined with another origin type.""" + if comid is not None: + raise ValueError( + "Provide only one origin type - comid cannot be provided with lat or long" + ) + if feature_source is not None or feature_id is not None: + raise ValueError( + "Provide only one origin type - feature_source and feature_id cannot" + " be provided with lat or long" + ) + + def _get_features_request( *, data_source: str | None, @@ -343,16 +360,7 @@ def _get_features_request( raise ValueError("Both lat and long are required") if lat is not None: - if comid is not None: - raise ValueError( - "Provide only one origin type - comid cannot be provided" - " with lat or long" - ) - if feature_source is not None or feature_id is not None: - raise ValueError( - "Provide only one origin type - feature_source and feature_id cannot" - " be provided with lat or long" - ) + _validate_lat_long_origin(comid, feature_source, feature_id) return f"{_api_base()}/comid/position", {"coords": f"POINT({long} {lat})"} if (comid is not None or data_source is not None) and navigation_mode is None: @@ -418,6 +426,37 @@ def get_features_by_data_source(data_source: str) -> gpd.GeoDataFrame: return gdf +def _search_basin(feature_source: str | None, feature_id: str | None) -> dict[str, Any]: + """Handle ``find='basin'`` for :func:`search`.""" + if feature_source is None or feature_id is None: + raise ValueError("feature_source and feature_id are required to find a basin") + return get_basin(feature_source=feature_source, feature_id=feature_id, as_json=True) + + +def _search_flowlines( + *, + navigation_mode: str | None, + distance: int, + feature_source: str | None, + feature_id: str | None, + comid: int | None, +) -> dict[str, Any]: + """Handle ``find='flowlines'`` for :func:`search`.""" + if navigation_mode is None: + raise ValueError( + "navigation_mode is required for find='flowlines';" + f" allowed values are {_VALID_NAVIGATION_MODES}" + ) + return get_flowlines( + navigation_mode=navigation_mode, + distance=distance, + feature_source=feature_source, + feature_id=feature_id, + comid=comid, + as_json=True, + ) + + def search( feature_source: str | None = None, feature_id: str | None = None, @@ -515,29 +554,18 @@ def search( return get_features(lat=lat, long=long, as_json=True) if find == "basin": - if feature_source is None or feature_id is None: - raise ValueError( - "feature_source and feature_id are required to find a basin" - ) - return get_basin( - feature_source=feature_source, feature_id=feature_id, as_json=True - ) + return _search_basin(feature_source, feature_id) if find == "flowlines": - if navigation_mode is None: - raise ValueError( - "navigation_mode is required for find='flowlines';" - f" allowed values are {_VALID_NAVIGATION_MODES}" - ) - return get_flowlines( + return _search_flowlines( navigation_mode=navigation_mode, distance=distance, feature_source=feature_source, feature_id=feature_id, comid=comid, - as_json=True, ) - # here find == 'features' + + # find == 'features' return get_features( data_source=data_source, navigation_mode=navigation_mode, diff --git a/dataretrieval/nwis.py b/dataretrieval/nwis.py index 8703e679..9221513e 100644 --- a/dataretrieval/nwis.py +++ b/dataretrieval/nwis.py @@ -127,6 +127,17 @@ def _parse_json_or_raise(response: httpx.Response) -> pd.DataFrame: raise +def _localize_datetime_index(df: pd.DataFrame) -> pd.DataFrame: + """Localize a naive datetime index (or multi-index level) to UTC.""" + if hasattr(df.index, "levels"): + # Multi-index: localize the datetime level (level 1) + if hasattr(df.index.levels[1], "tzinfo") and df.index.levels[1].tzinfo is None: + df = df.tz_localize("UTC", level=1) + elif hasattr(df.index, "tzinfo") and df.index.tzinfo is None: + df = df.tz_localize("UTC") + return df + + def format_response( df: pd.DataFrame, service: str | None = None, **kwargs: Any ) -> pd.DataFrame: @@ -162,22 +173,15 @@ def format_response( geoms = gpd.points_from_xy(df.dec_long_va.values, df.dec_lat_va.values) df = gpd.GeoDataFrame(df, geometry=geoms, crs=_CRS) - # check for multiple sites: if "datetime" not in df.columns: - # XXX: consider making site_no index return df - elif len(df["site_no"].unique()) > 1 and mi: - # setup multi-index + if len(df["site_no"].unique()) > 1 and mi: df.set_index(["site_no", "datetime"], inplace=True) - if hasattr(df.index.levels[1], "tzinfo") and df.index.levels[1].tzinfo is None: - df = df.tz_localize("UTC", level=1) - else: df.set_index(["datetime"], inplace=True) - if hasattr(df.index, "tzinfo") and df.index.tzinfo is None: - df = df.tz_localize("UTC") + df = _localize_datetime_index(df) return df.sort_index() @@ -1005,91 +1009,97 @@ def get_record( raise TypeError(f"{service} service not yet implemented") -def _read_json(json: dict[str, Any]) -> pd.DataFrame: - """Read a NWIS Water Services formatted JSON into a ``pandas.DataFrame``. +def _site_block_boundaries(site_list: list[str]) -> list[int]: + """Return indices where the site number changes, bookended by 0 and len. - Parameters - ---------- - json: dict - A JSON dictionary response to be parsed into a ``pandas.DataFrame``. + For example, given ``['A', 'A', 'B']`` returns ``[0, 2, 3]``. + """ + boundaries = [0] + boundaries.extend( + i + 1 + for i, (a, b) in enumerate(zip(site_list[:-1], site_list[1:], strict=False)) + if a != b + ) + boundaries.append(len(site_list)) + return boundaries - Returns - ------- - df: ``pandas.DataFrame`` - Time series data from the NWIS JSON. - """ - all_site_dfs = [] +def _build_column_name(param_cd: str, method: str, option: str | None) -> str: + """Derive the DataFrame column name for a parameter record.""" + col_name = param_cd + if method: + col_name = f"{col_name}_{method.strip('[]()').lower()}" + if option: + col_name = f"{col_name}_{option}" + return col_name - site_list = [ - ts["sourceInfo"]["siteCode"][0]["value"] for ts in json["value"]["timeSeries"] - ] - # create a list of indexes for each change in site no - # for example, [0, 21, 22] would be the first and last indices - index_list = [0] - index_list.extend( - [ - i + 1 - for i, (a, b) in enumerate(zip(site_list[:-1], site_list[1:], strict=False)) - if a != b - ] +def _parse_parameter_record( + record_json: list[dict[str, Any]], col_name: str +) -> pd.DataFrame: + """Parse a single parameter's value list into a renamed DataFrame.""" + record_df = pd.DataFrame(record_json) + record_df["value"] = pd.to_numeric(record_df["value"], errors="coerce") + record_df["qualifiers"] = ( + record_df["qualifiers"].astype(str).str.strip("[]").str.replace("'", "") ) - index_list.append(len(site_list)) - - for start, end in zip(index_list[:-1], index_list[1:], strict=False): - # grab a block containing timeseries 0:21, - # which are all from the same site - site_block = json["value"]["timeSeries"][start:end] - if not site_block: - continue + record_df.rename( + columns={ + "value": col_name, + "dateTime": "datetime", + "qualifiers": col_name + "_cd", + }, + inplace=True, + ) + return record_df - site_no = site_block[0]["sourceInfo"]["siteCode"][0]["value"] - site_df = pd.DataFrame(columns=["datetime"]) - for timeseries in site_block: - param_cd = timeseries["variable"]["variableCode"][0]["value"] - # check whether min, max, mean record XXX - option = timeseries["variable"]["options"]["option"][0].get("value") +def _parse_site_block(site_block: list[dict[str, Any]]) -> pd.DataFrame: + """Parse all timeseries in one site's block into a single DataFrame.""" + site_no = site_block[0]["sourceInfo"]["siteCode"][0]["value"] + site_df = pd.DataFrame(columns=["datetime"]) - for parameter in timeseries["values"]: - col_name = param_cd - method = parameter["method"][0]["methodDescription"] + for timeseries in site_block: + param_cd = timeseries["variable"]["variableCode"][0]["value"] + option = timeseries["variable"]["options"]["option"][0].get("value") - if method: - method = method.strip("[]()").lower() - col_name = f"{col_name}_{method}" + for parameter in timeseries["values"]: + method = parameter["method"][0]["methodDescription"] + col_name = _build_column_name(param_cd, method, option) + record_json = parameter["value"] + if not record_json: + continue + record_df = _parse_parameter_record(record_json, col_name) + site_df = site_df.merge(record_df, how="outer", on="datetime") - if option: - col_name = f"{col_name}_{option}" + site_df["site_no"] = site_no + return site_df - record_json = parameter["value"] - if not record_json: - continue +def _read_json(json: dict[str, Any]) -> pd.DataFrame: + """Read a NWIS Water Services formatted JSON into a ``pandas.DataFrame``. - record_df = pd.DataFrame(record_json) - record_df["value"] = pd.to_numeric(record_df["value"], errors="coerce") - record_df["qualifiers"] = ( - record_df["qualifiers"] - .astype(str) - .str.strip("[]") - .str.replace("'", "") - ) + Parameters + ---------- + json: dict + A JSON dictionary response to be parsed into a ``pandas.DataFrame``. - record_df.rename( - columns={ - "value": col_name, - "dateTime": "datetime", - "qualifiers": col_name + "_cd", - }, - inplace=True, - ) + Returns + ------- + df: ``pandas.DataFrame`` + Time series data from the NWIS JSON. - site_df = site_df.merge(record_df, how="outer", on="datetime") + """ + time_series = json["value"]["timeSeries"] + site_list = [ts["sourceInfo"]["siteCode"][0]["value"] for ts in time_series] + boundaries = _site_block_boundaries(site_list) - site_df["site_no"] = site_no - all_site_dfs.append(site_df) + all_site_dfs = [] + for start, end in zip(boundaries[:-1], boundaries[1:], strict=False): + site_block = time_series[start:end] + if not site_block: + continue + all_site_dfs.append(_parse_site_block(site_block)) if not all_site_dfs: return pd.DataFrame(columns=["site_no", "datetime"]) diff --git a/dataretrieval/ogc/dates.py b/dataretrieval/ogc/dates.py index 248e667c..f12ced9c 100644 --- a/dataretrieval/ogc/dates.py +++ b/dataretrieval/ogc/dates.py @@ -59,9 +59,14 @@ def _parse_datetime(value: str) -> datetime | None: return None +def _is_blank(dt: str | None) -> bool: + """True for a None, NaN, or empty-string element.""" + return dt is None or bool(pd.isna(dt)) or dt == "" + + def _format_one(dt: str | None, *, date: bool) -> str | None: """Format a single datetime element for inclusion in the API time arg.""" - if pd.isna(dt) or dt == "" or dt is None: + if dt is None or _is_blank(dt): return ".." parsed = _parse_datetime(dt) if parsed is None: @@ -76,6 +81,30 @@ def _format_one(dt: str | None, *, date: bool) -> str | None: return aware.astimezone(ZoneInfo("UTC")).strftime("%Y-%m-%dT%H:%M:%SZ") +def _coerce_to_list( + datetime_input: str | Sequence[str | None], +) -> list[str | None]: + """Normalize datetime input to a list, raising on invalid shapes.""" + if isinstance(datetime_input, str): + return [datetime_input] + if isinstance(datetime_input, Mapping): + raise TypeError( + f"date input must be a string or sequence of strings, " + f"not {type(datetime_input).__name__}." + ) + return list(datetime_input) + + +def _is_passthrough(single: str) -> bool: + """True when a single-element input should be returned as-is.""" + return bool(_DURATION_RE.match(single) or "/" in single) + + +def _all_blank(items: list[str | None]) -> bool: + """True when every element is None, NaN, or the empty string.""" + return all(_is_blank(dt) for dt in items) + + def _format_api_dates( datetime_input: str | Sequence[str | None] | None, date: bool = False ) -> str | None: @@ -128,37 +157,21 @@ def _format_api_dates( if datetime_input is None: return None - # Convert single string to list for uniform processing - if isinstance(datetime_input, str): - datetime_input = [datetime_input] - elif isinstance(datetime_input, Mapping): - # `list(mapping)` returns keys, which silently accepts the wrong shape. - raise TypeError( - f"date input must be a string or sequence of strings, " - f"not {type(datetime_input).__name__}." - ) - elif not isinstance(datetime_input, (list, tuple)): - # Materialize any other iterable (pandas.Series, numpy.ndarray, - # generator, ...) so the len()/subscript operations below work. - datetime_input = list(datetime_input) + items = _coerce_to_list(datetime_input) - # Check for null or all NA and return None - if all(pd.isna(dt) or dt == "" or dt is None for dt in datetime_input): + if _all_blank(items): return None - if len(datetime_input) > 2: + if len(items) > 2: raise ValueError("datetime_input should only include 1-2 values") # Pass through duration ("P7D", "PT36H") and pre-formatted interval ("a/b") - # strings untouched. - if len(datetime_input) == 1 and isinstance(datetime_input[0], str): - single = datetime_input[0] - if _DURATION_RE.match(single) or "/" in single: - return single + if len(items) == 1 and isinstance(items[0], str) and _is_passthrough(items[0]): + return items[0] # Format each element; any element that fails to parse invalidates the range. formatted: list[str] = [] - for dt in datetime_input: + for dt in items: one = _format_one(dt, date=date) if one is None: return None diff --git a/dataretrieval/ogc/errors.py b/dataretrieval/ogc/errors.py index 597f86e8..c39fa283 100644 --- a/dataretrieval/ogc/errors.py +++ b/dataretrieval/ogc/errors.py @@ -66,6 +66,14 @@ def _error_body(resp: httpx.Response) -> str: ) +def _clean_field(value: object | None) -> str | None: + """Normalize a JSON error field: strip whitespace and trailing dot.""" + if value is None: + return None + text = str(value).strip().rstrip(".") + return text or None + + def _json_error_detail(resp: httpx.Response) -> str | None: """Render a supported JSON error body, or ``None`` for another shape.""" try: @@ -79,15 +87,11 @@ def _json_error_detail(resp: httpx.Response) -> str | None: if not isinstance(candidate, dict): candidate = body - def clean(value: object | None) -> str | None: - if value is None: - return None - text = str(value).strip().rstrip(".") - return text or None - - code = clean(candidate.get("code")) - detail = clean(candidate.get("description")) or clean(candidate.get("message")) - parts = [part for part in (code, detail) if part is not None] + code = _clean_field(candidate.get("code")) + detail = _clean_field(candidate.get("description")) or _clean_field( + candidate.get("message") + ) + parts = [p for p in (code, detail) if p] return ". ".join(parts) + "." if parts else None diff --git a/dataretrieval/ogc/filters.py b/dataretrieval/ogc/filters.py index 50875409..2d4bc451 100644 --- a/dataretrieval/ogc/filters.py +++ b/dataretrieval/ogc/filters.py @@ -20,6 +20,7 @@ from __future__ import annotations import re +from collections.abc import Iterator from typing import Any, Literal FILTER_LANG = Literal["cql-text", "cql-json"] @@ -90,41 +91,104 @@ def _resume_after_or(expr: str, i: int) -> int | None: return _skip_space(expr, after_word) -def _split_top_level_or(expr: str) -> list[str]: - """Split ``expr`` at each top-level ``OR``, respecting quotes and parens. +def _skip_quoted(expr: str, i: int) -> int: + """Index just past the quoted span opening at ``i``. - ``OR`` tokens inside ``(A OR B)`` or ``'word OR word'`` are left alone. - Matching is case-insensitive; whitespace around each part is stripped; - empty parts are dropped. + An unterminated quote swallows the rest of the expression. A doubled + ``''`` escape reads as close-then-reopen, which nets to the same state. + """ + close = expr.find(expr[i], i + 1) + return len(expr) if close == -1 else close + 1 + + +def _iter_top_level_spaces(expr: str) -> Iterator[int]: + """Yield the indices of whitespace outside quotes and parens. + + Owns the depth/quote state so callers only see split candidates; quoted + spans are skipped wholesale. """ - parts: list[str] = [] - last = 0 depth = 0 - in_quote: str | None = None i = 0 n = len(expr) while i < n: ch = expr[i] - if in_quote is not None: - if ch == in_quote: - in_quote = None - elif ch in ("'", '"'): - in_quote = ch - elif ch == "(": + if ch in ("'", '"'): + i = _skip_quoted(expr, i) + continue + if ch == "(": depth += 1 elif ch == ")": depth -= 1 elif depth == 0 and ch.isspace(): - resume = _resume_after_or(expr, i + 1) - if resume is not None: - parts.append(expr[last:i].strip()) - last = i = resume - continue + yield i i += 1 + + +def _split_top_level_or(expr: str) -> list[str]: + """Split ``expr`` at each top-level ``OR``, respecting quotes and parens. + + ``OR`` tokens inside ``(A OR B)`` or ``'word OR word'`` are left alone. + Matching is case-insensitive; whitespace around each part is stripped; + empty parts are dropped. + """ + parts: list[str] = [] + last = 0 + for i in _iter_top_level_spaces(expr): + if i < last: + continue # inside an already-consumed OR separator + resume = _resume_after_or(expr, i + 1) + if resume is not None: + parts.append(expr[last:i].strip()) + last = resume parts.append(expr[last:].strip()) return [p for p in parts if p] +def _numeric_pitfall_error(field: str, offense: str) -> ValueError: + """Build the error for an unquoted numeric comparison.""" + return ValueError( + f"Filter uses an unquoted numeric comparison against {field!r} " + f"(``{offense}``). Every queryable on the Water Data API is " + f"typed as a string, so the server rejects unquoted numeric " + f"literals with HTTP 500; even quoting the literal gives a " + f"lexicographic comparison (``value > '10'`` matches " + f"``value='34.52'``, ``parameter_code = '60'`` matches nothing " + f"because the real codes are ``'00060'``-shaped). For a true " + f"numeric filter, fetch a wider result and reduce in pandas." + ) + + +def _check_compare(masked: str) -> None: + """Raise on a bare ``field op number`` or ``number op field`` pattern.""" + compare = _NUMERIC_COMPARE_RE.search(masked) + if not compare: + return + field = compare.group("field1") or compare.group("field2") + op = compare.group("op1") or compare.group("op2") + num = compare.group("num1") or compare.group("num2") + raise _numeric_pitfall_error(field, f"{field} {op} {num}") + + +def _check_in_membership(masked: str) -> None: + """Raise on a ``field [NOT] IN (…numeric…)`` pattern.""" + membership = _IN_NUMERIC_RE.search(masked) + if not membership: + return + field = membership.group("field") + op = "NOT IN" if membership.group("negated") else "IN" + raise _numeric_pitfall_error(field, f"{field} {op} (…)") + + +def _check_between(masked: str) -> None: + """Raise on a ``field [NOT] BETWEEN … AND …`` pattern with numerics.""" + between = _BETWEEN_NUMERIC_RE.search(masked) + if not between: + return + field = between.group("field") + op = "NOT BETWEEN" if between.group("negated") else "BETWEEN" + raise _numeric_pitfall_error(field, f"{field} {op} …") + + def _check_numeric_filter_pitfall(filter_expr: str) -> None: """Raise if the filter pairs a field with an unquoted numeric literal. @@ -145,37 +209,9 @@ def _check_numeric_filter_pitfall(filter_expr: str) -> None: masked = ( _QUOTED_STR_RE.sub("''", filter_expr) if "'" in filter_expr else filter_expr ) - - def fail(field: str, offense: str) -> None: - raise ValueError( - f"Filter uses an unquoted numeric comparison against {field!r} " - f"(``{offense}``). Every queryable on the Water Data API is " - f"typed as a string, so the server rejects unquoted numeric " - f"literals with HTTP 500; even quoting the literal gives a " - f"lexicographic comparison (``value > '10'`` matches " - f"``value='34.52'``, ``parameter_code = '60'`` matches nothing " - f"because the real codes are ``'00060'``-shaped). For a true " - f"numeric filter, fetch a wider result and reduce in pandas." - ) - - compare = _NUMERIC_COMPARE_RE.search(masked) - if compare: - field = compare.group("field1") or compare.group("field2") - op = compare.group("op1") or compare.group("op2") - num = compare.group("num1") or compare.group("num2") - fail(field, f"{field} {op} {num}") - - membership = _IN_NUMERIC_RE.search(masked) - if membership: - field = membership.group("field") - op = "NOT IN" if membership.group("negated") else "IN" - fail(field, f"{field} {op} (…)") - - between = _BETWEEN_NUMERIC_RE.search(masked) - if between: - field = between.group("field") - op = "NOT BETWEEN" if between.group("negated") else "BETWEEN" - fail(field, f"{field} {op} …") + _check_compare(masked) + _check_in_membership(masked) + _check_between(masked) def _is_chunkable(filter_expr: Any, filter_lang: Any) -> bool: diff --git a/dataretrieval/ogc/planning.py b/dataretrieval/ogc/planning.py index e5af5411..ba208a51 100644 --- a/dataretrieval/ogc/planning.py +++ b/dataretrieval/ogc/planning.py @@ -92,6 +92,36 @@ def _request_bytes(req: httpx.Request) -> int: return len(str(req.url)) + len(req.content) +def _try_build( + build_request: Callable[..., httpx.Request], + args: dict[str, Any], +) -> httpx.Request | None: + """Attempt to construct a request, returning ``None`` on overflow. + + ``httpx.URL`` enforces a hard 64 KB cap per URL component and raises + ``httpx.InvalidURL`` for anything bigger. Both :func:`_safe_request_bytes` + and :class:`ChunkPlan`'s initial-request probe need exactly this + "build-or-None" step, so it lives here once. + + Parameters + ---------- + build_request : Callable[..., httpx.Request] + Factory that turns a kwargs dict into a sized request. + args : dict[str, Any] + Per-chunk kwargs to pass through to ``build_request``. + + Returns + ------- + httpx.Request or None + The built request, or ``None`` when construction raised + ``httpx.InvalidURL``. + """ + try: + return build_request(**args) + except httpx.InvalidURL: + return None + + def _safe_request_bytes( build_request: Callable[..., httpx.Request], args: dict[str, Any], @@ -122,11 +152,36 @@ def _safe_request_bytes( ``url_limit + 1`` so the planner's "too large" branch keeps halving. """ - try: - req = build_request(**args) - except httpx.InvalidURL: - return url_limit + 1 - return _request_bytes(req) + req = _try_build(build_request, args) + return _request_bytes(req) if req is not None else url_limit + 1 + + +def _check_unchunkable_request( + args: dict[str, Any], + build_request: Callable[..., httpx.Request], + url_limit: int, +) -> None: + """Enforce the byte budget on a request with no chunkable axis. + + Passthrough when the single request fits or when the filter is in a + language the chunker doesn't manage (cql-json) — the server, not us, + judges that one. Raises + :class:`~dataretrieval.exceptions.Unchunkable` when the request is + over budget and has nothing to split. + """ + if _safe_request_bytes(build_request, args, url_limit) <= url_limit: + return + filter_expr = args.get("filter") + if filter_expr is not None and not _is_chunkable( + filter_expr, args.get("filter_lang") + ): + return + raise Unchunkable( + f"Request exceeds {url_limit} bytes (URL + body) and has no " + f"chunkable multi-value argument to split (e.g. a single large " + f"CQL `IN` clause, or one oversized value). Narrow the query, " + f"simplify the filter, or split the call manually." + ) @dataclass(frozen=True) @@ -201,6 +256,22 @@ def render(self, chunk: list[str]) -> Any: return list(chunk) if self.joiner == _LIST_SEP else self.joiner.join(chunk) +def _filter_axis(args: dict[str, Any]) -> _Axis | None: + """Build the filter axis from CQL-text ``filter``, if chunkable. + + Returns an :class:`_Axis` whose atoms are top-level OR-clauses when the + filter has two or more splittable clauses; ``None`` otherwise. + """ + filter_expr = args.get("filter") + if filter_expr is None or not _is_chunkable(filter_expr, args.get("filter_lang")): + return None + _check_numeric_filter_pitfall(filter_expr) + clauses = _split_top_level_or(filter_expr) + if len(clauses) < 2: + return None + return _Axis(arg_key="filter", atoms=tuple(clauses), joiner=_OR_SEP) + + def _extract_axes(args: dict[str, Any]) -> list[_Axis]: """ Build the chunkable-axis set from a request's args. @@ -224,19 +295,16 @@ def _extract_axes(args: dict[str, Any]) -> list[_Axis]: per eligible kwarg, in ``args`` order), then the filter axis if present. """ - axes: list[_Axis] = [] - for key, value in args.items(): - if key in _NEVER_CHUNK: - continue - if isinstance(value, (list, tuple)) and len(value) > 1: - axes.append(_Axis(arg_key=key, atoms=tuple(value), joiner=_LIST_SEP)) - - filter_expr = args.get("filter") - if filter_expr is not None and _is_chunkable(filter_expr, args.get("filter_lang")): - _check_numeric_filter_pitfall(filter_expr) - clauses = _split_top_level_or(filter_expr) - if len(clauses) >= 2: - axes.append(_Axis(arg_key="filter", atoms=tuple(clauses), joiner=_OR_SEP)) + axes: list[_Axis] = [ + _Axis(arg_key=key, atoms=tuple(value), joiner=_LIST_SEP) + for key, value in args.items() + if key not in _NEVER_CHUNK + and isinstance(value, (list, tuple)) + and len(value) > 1 + ] + fax = _filter_axis(args) + if fax is not None: + axes.append(fax) return axes @@ -330,12 +398,6 @@ def __init__( max_chunks: int = 1, ) -> None: if max_chunks < 1: - # ``max_chunks`` is a chunk *count*: the minimum is ``1`` - # (the ambient default outside any ``parallel_chunks`` block), - # which means "off — no extra fan-out". ``0`` or negative is a - # meaningless count and can only be a caller bug, so fail loudly - # rather than silently no-op. The public ``parallel_chunks(n)`` - # already rejects ``n < 1``; this guards direct construction. raise ValueError( f"max_chunks must be >= 1 (1 disables fan-out); got {max_chunks!r}." ) @@ -347,73 +409,27 @@ def __init__( axes = _extract_axes(args) if not axes: - # No chunkable axis: nothing to split, and ``parallel_chunks`` has - # nothing to act on either. If the single request fits, run it - # verbatim (the common passthrough). ``_safe_request_bytes`` treats - # an un-constructable URL (httpx.InvalidURL, > 64 KB) as over budget. - if _safe_request_bytes(build_request, args, url_limit) <= url_limit: - return - # Over budget. A filter the chunker doesn't manage — cql-json — is - # passed through unchanged (chunking applies only to cql-text); the - # server, not us, judges it. Otherwise this is an in-domain shape we - # would normally chunk but can't: a single large CQL ``IN`` clause - # with no top-level ``OR``, or one oversized value. Raise an - # actionable error instead of shipping it for an opaque HTTP 414. - filter_expr = args.get("filter") - if filter_expr is not None and not _is_chunkable( - filter_expr, args.get("filter_lang") - ): - return - raise Unchunkable( - f"Request exceeds {url_limit} bytes (URL + body) and has no " - f"chunkable multi-value argument to split (e.g. a single large " - f"CQL `IN` clause, or one oversized value). Narrow the query, " - f"simplify the filter, or split the call manually." - ) - - # Constructing the initial request can itself trip - # ``httpx.InvalidURL`` (URL > 64 KB) — that's the canonical - # "needs chunking" signal, so swallow it and proceed to plan. - # When the unchunked URL does build, preserve it as ``canonical_url`` - # so ``BaseMetadata.url`` echoes the user's original query verbatim. - # Only fall back to a worst-case chunk URL when the URL itself - # can't be constructed. - try: - initial_request = build_request(**args) - except httpx.InvalidURL: - initial_request = None + _check_unchunkable_request(args, build_request, url_limit) + return + # When the un-chunked URL builds, preserve it as ``canonical_url`` so + # ``BaseMetadata.url`` echoes the user's original query verbatim. + initial_request = _try_build(build_request, args) fits = False if initial_request is not None: self.canonical_url = str(initial_request.url) fits = _request_bytes(initial_request) <= url_limit - # A request that already fits and hasn't opted into finer chunking is - # the common passthrough: leave ``axes``/``chunks`` empty so - # ``total == 1`` and ``iter_chunk_args`` yields the original args - # verbatim. ``max_chunks == 1`` (off / no extra fan-out) means - # "don't split", so it takes this path; only ``max_chunks >= 2`` asks - # for extra fan-out and sets the axes up to be refined below. if fits and max_chunks <= 1: return self.axes = axes self.chunks = {axis.arg_key: [list(axis.atoms)] for axis in axes} if not fits: - # Hard pass: greedy-halve until every worst-case chunk fits - # the byte budget (may raise ``Unchunkable``). self._plan(build_request, url_limit) - # Soft pass: optionally split further than the byte budget requires. - # Purely additive — never re-raises, and the byte budget stays - # satisfied; a no-op at ``max_chunks == 1``. self._refine(max_chunks) if self.canonical_url is None: - # Original URL was un-constructable (httpx.InvalidURL); fall - # back to the worst-case chunk URL so - # ``BaseMetadata.url`` still surfaces something - # informative. If even that overflows, leave canonical_url - # as None (set above) and let the response's own URL stand. with suppress(httpx.InvalidURL): self.canonical_url = str(build_request(**self._worst_case_args()).url) @@ -441,17 +457,7 @@ def _plan( if _safe_request_bytes(build_request, worst, url_limit) <= url_limit: return - biggest_axis: _Axis | None = None - biggest_idx = -1 - biggest_size = -1 - for axis in self.axes: - for idx, chunk in enumerate(self.chunks[axis.arg_key]): - if len(chunk) <= 1: - continue - size = axis.chunk_bytes(chunk) - if size > biggest_size: - biggest_axis, biggest_idx, biggest_size = axis, idx, size - + biggest_axis, biggest_idx = self._largest_splittable_chunk_by_bytes() if biggest_axis is None: raise Unchunkable( f"Request exceeds {url_limit} bytes (URL + body) at the " @@ -461,6 +467,23 @@ def _plan( ) _split_at(self.chunks[biggest_axis.arg_key], biggest_idx) + def _largest_splittable_chunk_by_bytes(self) -> tuple[_Axis | None, int]: + """Find the largest splittable chunk ranked by URL-encoded byte size. + + Returns ``(axis, index)`` of the biggest chunk with more than one atom, + or ``(None, -1)`` when every axis is at one atom per chunk (saturated). + """ + biggest_axis: _Axis | None = None + biggest_idx = -1 + biggest_size = -1 + for axis in self.axes: + idx, size = self._largest_chunk_in( + self.chunks[axis.arg_key], key=axis.chunk_bytes + ) + if size > biggest_size: + biggest_axis, biggest_idx, biggest_size = axis, idx, size + return biggest_axis, biggest_idx + def _refine(self, max_chunks: int) -> None: """ Fan the plan out more finely than the byte budget alone requires. @@ -496,34 +519,58 @@ def _refine(self, max_chunks: int) -> None: total = self.total if total >= max_chunks: return - # Largest splittable chunk among the axes whose split still fits the - # cap. Splitting any chunk of an axis with ``k`` chunks turns that - # ``k`` into ``k+1``, so it adds ``total // k`` chunks (the - # product of the other axes) regardless of which chunk. Hence the - # budget test is per axis, not per chunk. Skipping an over-budget - # axis makes ``max_chunks`` a true ceiling. The ranking key is atom - # count (``len``), not URL bytes like ``_plan`` — this pass balances - # work across chunks rather than fitting a byte budget. A - # chunk of size 1 can't be split further. Stable input order breaks - # ties by axis order, then lowest index within an axis. - candidate: tuple[_Axis, int] | None = None - candidate_size = -1 - for axis in self.axes: - axis_chunks = self.chunks[axis.arg_key] - if total + total // len(axis_chunks) > max_chunks: - continue # any split of this axis would overshoot the cap - for idx, chunk in enumerate(axis_chunks): - if len(chunk) <= 1: - continue - if len(chunk) > candidate_size: - candidate, candidate_size = (axis, idx), len(chunk) + candidate = self._best_refine_candidate(total, max_chunks) if candidate is None: - # Every axis is saturated at one atom per chunk or would - # overshoot the cap; stop below it rather than exceed it. return axis, idx = candidate _split_at(self.chunks[axis.arg_key], idx) + @staticmethod + def _largest_chunk_in( + axis_chunks: list[list[str]], + key: Callable[[list[str]], int] = len, + ) -> tuple[int, int]: + """Return ``(index, key(chunk))`` of the largest splittable chunk. + + A chunk is splittable when it has more than one atom; *key* ranks the + qualifying chunks (atom count by default, URL bytes for the byte + pass). Returns ``(-1, -1)`` when no chunk qualifies. + """ + best_idx = -1 + best_size = -1 + for idx, chunk in enumerate(axis_chunks): + if len(chunk) <= 1: + continue + size = key(chunk) + if size > best_size: + best_idx, best_size = idx, size + return best_idx, best_size + + def _best_refine_candidate( + self, total: int, max_chunks: int + ) -> tuple[_Axis, int] | None: + """Find the best chunk to split during the refine pass. + + Returns the largest splittable chunk (by atom count) among axes whose + split stays within the ``max_chunks`` cap, or ``None`` when no + in-budget split remains. Splitting any chunk of an axis with ``k`` + chunks adds ``total // k`` chunks (the product of the other axes), + so the budget test is per axis rather than per chunk. The ranking key + is atom count (not URL bytes like ``_plan``) because this pass + balances work across chunks rather than fitting a byte budget. + Stable input order breaks ties by axis order, then lowest index. + """ + candidate: tuple[_Axis, int] | None = None + candidate_size = -1 + for axis in self.axes: + axis_chunks = self.chunks[axis.arg_key] + if total + total // len(axis_chunks) > max_chunks: + continue # any split of this axis would overshoot the cap + axis_best, axis_best_size = self._largest_chunk_in(axis_chunks) + if axis_best_size > candidate_size: + candidate, candidate_size = (axis, axis_best), axis_best_size + return candidate + def _worst_case_args(self) -> dict[str, Any]: """ Args for the largest chunk the current partition will issue. diff --git a/dataretrieval/ogc/requests.py b/dataretrieval/ogc/requests.py index d3b76ae4..c82752cf 100644 --- a/dataretrieval/ogc/requests.py +++ b/dataretrieval/ogc/requests.py @@ -95,31 +95,46 @@ def _ogc_query_params( return params -def _partition_request_params( - params: dict[str, Any], *, use_cql2: bool +def _is_post_param(value: Any) -> bool: + """True when ``value`` is a multi-value list suitable for CQL2 POST.""" + return isinstance(value, (list, tuple)) and len(value) > 1 + + +def _partition_cql2( + params: dict[str, Any], ) -> tuple[dict[str, Any], dict[str, Any]]: - """Split URL parameters from multi-value CQL2 POST predicates.""" - if use_cql2: - post_params = { - key: value - for key, value in params.items() - if isinstance(value, (list, tuple)) and len(value) > 1 - } - return ( - {key: value for key, value in params.items() if key not in post_params}, - post_params, - ) + """CQL2 path: multi-value lists go to POST body, rest stay as URL params.""" + post_params = {key: value for key, value in params.items() if _is_post_param(value)} + url_params = {key: value for key, value in params.items() if key not in post_params} + return url_params, post_params + + +def _join_get_value(value: Any) -> Any: + """Comma-join list/tuple values for GET params; pass scalars through.""" + if isinstance(value, (list, tuple)): + return ",".join(str(item) for item in value) + return value + +def _partition_get(params: dict[str, Any]) -> tuple[dict[str, Any], dict[str, Any]]: + """GET path: comma-join multi-value lists, drop empty lists.""" get_params = { - key: ",".join(str(item) for item in value) - if isinstance(value, (list, tuple)) - else value + key: _join_get_value(value) for key, value in params.items() if not (isinstance(value, (list, tuple)) and len(value) == 0) } return get_params, {} +def _partition_request_params( + params: dict[str, Any], *, use_cql2: bool +) -> tuple[dict[str, Any], dict[str, Any]]: + """Split URL parameters from multi-value CQL2 POST predicates.""" + if use_cql2: + return _partition_cql2(params) + return _partition_get(params) + + def _items_url(collection: str, base_url: str) -> str: """The OGC items endpoint for ``collection`` under ``base_url``.""" return f"{base_url}/collections/{collection}/items" @@ -292,6 +307,27 @@ def _check_monitoring_location_id( return value +def _normalize_arg( + key: str, + value: Any, + no_normalize: frozenset[str], +) -> Any: + """Normalize a single request argument value based on its key.""" + if key == "monitoring_location_id": + return _check_monitoring_location_id(value) + if key == "properties": + return _as_str_list(value, key) + if ( + key in no_normalize + and isinstance(value, Iterable) + and not isinstance(value, str) + ): + return value.tolist() if hasattr(value, "tolist") else list(value) + if isinstance(value, str) or not isinstance(value, Iterable): + return value + return _normalize_str_iterable(value, key) + + def prepare_request_args( local_vars: dict[str, Any], exclude: set[str] | None = None, @@ -310,9 +346,6 @@ def prepare_request_args( by forgetting to union them back in. """ no_normalize = _NO_NORMALIZE_PARAMS | frozenset(extra_no_normalize) - # Both spellings: this drops the caller's collection selector out of the - # query string, and public getters still name that local ``service`` - # (waterdata.get_samples) or ``service=`` during the get_cql deprecation. to_exclude = {"collection", "service", "output_id"} if exclude: to_exclude.update(exclude) @@ -321,14 +354,5 @@ def prepare_request_args( for k, v in local_vars.items(): if k in to_exclude or v is None: continue - if k == "monitoring_location_id": - args[k] = _check_monitoring_location_id(v) - elif k == "properties": - args[k] = _as_str_list(v, k) - elif k in no_normalize and isinstance(v, Iterable) and not isinstance(v, str): - args[k] = v.tolist() if hasattr(v, "tolist") else list(v) - elif isinstance(v, str) or not isinstance(v, Iterable): - args[k] = v - else: - args[k] = _normalize_str_iterable(v, k) + args[k] = _normalize_arg(k, v, no_normalize) return args diff --git a/dataretrieval/ogc/shaping.py b/dataretrieval/ogc/shaping.py index 760e6f41..eb8a692b 100644 --- a/dataretrieval/ogc/shaping.py +++ b/dataretrieval/ogc/shaping.py @@ -58,18 +58,18 @@ def _empty_feature_frame( collection schema; page-level empties intentionally remain schema-light. """ result_columns = list(columns or []) - if include_geometry: - if "geometry" not in result_columns: - result_columns.append("geometry") - else: + if not include_geometry: result_columns = [name for name in result_columns if name != "geometry"] + elif "geometry" not in result_columns: + result_columns.append("geometry") data = {name: pd.Series(dtype=object) for name in result_columns} - if not geopd: - return pd.DataFrame(data, columns=result_columns) - - data["geometry"] = gpd.GeoSeries([], crs=_CRS) - return gpd.GeoDataFrame(data, columns=result_columns, geometry="geometry", crs=_CRS) + if geopd: + data["geometry"] = gpd.GeoSeries([], crs=_CRS) + return gpd.GeoDataFrame( + data, columns=result_columns, geometry="geometry", crs=_CRS + ) + return pd.DataFrame(data, columns=result_columns) def _attach_coordinates(df: pd.DataFrame, features: list[dict[str, Any]]) -> None: @@ -99,6 +99,25 @@ def _geo_feature_frame(features: list[dict[str, Any]]) -> pd.DataFrame: ) +def _plain_feature_frame( + features: list[dict[str, Any]], *, include_geometry: bool +) -> pd.DataFrame: + """Build a plain DataFrame from GeoJSON features.""" + properties = [feature.get("properties") or {} for feature in features] + df = pd.json_normalize(properties, sep="_") + df["id"] = [feature.get("id") for feature in features] + if include_geometry: + _attach_coordinates(df, features) + return df + + +def _spatial_feature_frame(features: list[dict[str, Any]]) -> pd.DataFrame: + """Build a GeoDataFrame from GeoJSON features with ``id`` first.""" + df = _geo_feature_frame(features) + df["id"] = [f.get("id") for f in features] + return df[["id"] + [col for col in df.columns if col != "id"]] + + def _get_resp_data( resp: httpx.Response, geopd: bool, @@ -160,23 +179,9 @@ def _get_resp_data( return _empty_feature_frame(geopd, include_geometry=include_geometry) if not geopd: - properties = [feature.get("properties") or {} for feature in features] - df = pd.json_normalize(properties, sep="_") - # Always materialize the feature-level ID (possibly all-None) so - # ``_arrange_cols`` can perform the documented collection-specific rename. - df["id"] = [feature.get("id") for feature in features] - if include_geometry: - _attach_coordinates(df, features) - return df + return _plain_feature_frame(features, include_geometry=include_geometry) - # A spatial request remains geospatial even when this particular page - # carries only null/missing geometries; changing frame family based on page - # contents makes pagination concat order-dependent. - df = _geo_feature_frame(features) - # Mirror the plain branch's defensive ``f.get("id")`` so a feature missing - # a top-level id yields None rather than a KeyError. - df["id"] = [f.get("id") for f in features] - return df[["id"] + [col for col in df.columns if col != "id"]] + return _spatial_feature_frame(features) def _deal_with_empty( @@ -246,31 +251,26 @@ def _arrange_cols( # Rename id column to output_id df = df.rename(columns={"id": output_id}) - if properties and not all(pd.isna(properties)): - # Don't alias the caller's list — we mutate below. - local_properties = list(properties) - if "geometry" in df.columns and "geometry" not in local_properties: - local_properties.append("geometry") - # 'id' is a valid collection column, but expose it under the - # collection-specific output_id name instead. - if "id" in local_properties: - local_properties[local_properties.index("id")] = output_id - df = df.loc[:, [col for col in local_properties if col in df.columns]] - - # Move meaningless-to-user, extra id columns to the end - # of the dataframe, if they exist - extra_id_col = set(df.columns).intersection(extra_id_cols) - - # If the arbitrary id column is returned (either due to properties - # being none or NaN), then move it to the end of the dataframe, but - # if part of properties, keep in requested order - if extra_id_col and (properties is None or all(pd.isna(properties))): - id_col_order = [col for col in df.columns if col not in extra_id_col] + list( - extra_id_col - ) - df = df.loc[:, id_col_order] + # --- No explicit properties: move meaningless extra-id cols to end --- + if not properties or all(pd.isna(properties)): + extra_id_col = set(df.columns).intersection(extra_id_cols) + if extra_id_col: + id_col_order = [ + col for col in df.columns if col not in extra_id_col + ] + list(extra_id_col) + df = df.loc[:, id_col_order] + return df - return df + # --- Explicit properties: select and reorder columns per the list --- + # Don't alias the caller's list — we mutate below. + local_properties = list(properties) + if "geometry" in df.columns and "geometry" not in local_properties: + local_properties.append("geometry") + # 'id' is a valid collection column, but expose it under the + # collection-specific output_id name instead. + if "id" in local_properties: + local_properties[local_properties.index("id")] = output_id + return df.loc[:, [col for col in local_properties if col in df.columns]] def _type_cols(df: pd.DataFrame, dialect: OgcDialect) -> pd.DataFrame: diff --git a/dataretrieval/transport/fanout.py b/dataretrieval/transport/fanout.py index 5f9d7646..d812f775 100644 --- a/dataretrieval/transport/fanout.py +++ b/dataretrieval/transport/fanout.py @@ -525,6 +525,51 @@ def resume(self) -> tuple[pd.DataFrame, Any]: portal.call(functools.partial(self._run, concurrency)), ) + def _handle_gather_failures( + self, results: list[tuple[pd.DataFrame, httpx.Response] | BaseException] + ) -> None: + """Apply failure-precedence rules to gather results. + + Failure precedence, in order: + + 1. Cancellation / interrupt signals (``CancelledError``, + ``KeyboardInterrupt``, ``SystemExit`` — non-``Exception``) + propagate unmodified; wrapping them as a transient would swallow + the user's stop signal. + 2. A non-transient failure (a real bug — unrecognized by + ``wrap_failure``) surfaces raw, so it isn't masked behind a + resumable handle for a transient sibling that landed later. + 3. Only when every failure is a recognized transient do we raise + the first as a resumable ``FanOutInterrupted``. + + ``wrap_failure`` is asked only for the one failure that is raised. + Asking it per failure would snapshot the combined frame N times (a + full concat over every completed chunk) and discard all but one, + which a batch of chunks failing together makes routine. + + Raises + ------ + FanOutInterrupted + When all failures are recognized transients. + BaseException + Non-``Exception`` signals or non-transient failures. + """ + failures = [r for r in results if isinstance(r, BaseException)] + for exc in failures: + if not isinstance(exc, Exception): + raise exc + for exc in failures: + if _classify_chunk_error(exc) is None: + raise self._normalize_failure(exc) + if not failures: + return + first_transient = failures[0] + interrupted = self.wrap_failure(first_transient) + if interrupted is None: + # Unreachable: classified as transient just above. + raise self._normalize_failure(first_transient) + raise interrupted from first_transient + async def _run(self, max_concurrent: int | None) -> tuple[pd.DataFrame, Any]: """ Gather every pending chunk over one shared @@ -623,44 +668,11 @@ async def track( # semaphore (held by ``_retry`` per attempt) is the only throttle. # ``return_exceptions`` keeps completed pairs after a sibling # fails, so partial state stays recoverable via :meth:`resume`. - # Failure precedence, in order: - # 1. Cancellation / interrupt signals (CancelledError, - # KeyboardInterrupt, SystemExit — non-Exception) propagate - # unmodified; wrapping them as a transient would swallow - # the user's stop signal. - # 2. A non-transient failure (a real bug — unrecognized by - # ``wrap_failure``) surfaces raw, so it isn't masked behind - # a resumable handle for a transient sibling that landed - # later. - # 3. Only when every failure is a recognized transient do we - # raise the first as a resumable ``FanOutInterrupted``. results = await asyncio.gather( *(track(index, item) for index, item in self._pending()), return_exceptions=True, ) - failures = [r for r in results if isinstance(r, BaseException)] - for exc in failures: - if not isinstance(exc, Exception): - raise exc - # Classify first, build once. Every failure has to be - # examined -- a non-transient sibling must surface raw -- but - # only the first transient is ever raised. Asking - # ``wrap_failure`` per failure would snapshot the combined - # frame N times (a full concat over every completed - # chunk) and discard all but one, which a batch of - # chunks failing together makes routine. - first_transient: BaseException | None = None - for exc in failures: - if _classify_chunk_error(exc) is None: - raise self._normalize_failure(exc) - if first_transient is None: - first_transient = exc - if first_transient is not None: - interrupted = self.wrap_failure(first_transient) - if interrupted is None: - # Unreachable: classified as transient just above. - raise self._normalize_failure(first_transient) - raise interrupted from first_transient + self._handle_gather_failures(results) return self.finalize(*self._combine_raw()) diff --git a/dataretrieval/waterdata/ratings.py b/dataretrieval/waterdata/ratings.py index d5aecd4a..b327bbdf 100644 --- a/dataretrieval/waterdata/ratings.py +++ b/dataretrieval/waterdata/ratings.py @@ -174,18 +174,8 @@ def get_ratings( """ monitoring_location_id = _check_monitoring_location_id(monitoring_location_id) file_types = _as_list(file_type) - invalid = [ft for ft in file_types if ft not in _VALID_FILE_TYPES] - if invalid: - raise ValueError( - f"Invalid file_type {invalid!r}; " - f"valid options are {list(_VALID_FILE_TYPES)}." - ) - - if time is not None and any(_DURATION_RE.match(str(v)) for v in _as_list(time)): - raise ValueError( - "ISO 8601 durations (e.g. 'P7D') are not supported in `time` " - "for the rating-curve service. Provide a date or interval instead." - ) + _validate_file_types(file_types) + _validate_time_no_duration(time) time_str = _format_api_dates(time) if time is not None else None # Mirror R: pin file_type server-side only when one type is requested. @@ -197,10 +187,7 @@ def get_ratings( if not download_and_parse: return features - requested = set(file_types) - matching = [ - f for f in features if f.get("properties", {}).get("file_type") in requested - ] + matching = _filter_features_by_type(features, file_types) if file_path is not None: os.makedirs(file_path, exist_ok=True) @@ -213,6 +200,37 @@ def _as_list(x: str | Iterable[str]) -> list[str]: return [x] if isinstance(x, str) else list(x) +def _validate_file_types(file_types: list[str]) -> None: + """Raise ValueError for any unrecognized file type.""" + invalid = [ft for ft in file_types if ft not in _VALID_FILE_TYPES] + if invalid: + raise ValueError( + f"Invalid file_type {invalid!r}; " + f"valid options are {list(_VALID_FILE_TYPES)}." + ) + + +def _validate_time_no_duration(time: str | list[str] | None) -> None: + """Raise ValueError if ``time`` contains an ISO 8601 duration.""" + if time is None: + return + if any(_DURATION_RE.match(str(v)) for v in _as_list(time)): + raise ValueError( + "ISO 8601 durations (e.g. 'P7D') are not supported in `time` " + "for the rating-curve service. Provide a date or interval instead." + ) + + +def _filter_features_by_type( + features: list[dict[str, Any]], file_types: list[str] +) -> list[dict[str, Any]]: + """Return only features whose file_type matches the requested types.""" + requested = set(file_types) + return [ + f for f in features if f.get("properties", {}).get("file_type") in requested + ] + + def _build_filter( monitoring_location_id: str | list[str] | None, file_type: str | None, diff --git a/dataretrieval/waterdata/stats.py b/dataretrieval/waterdata/stats.py index 09ee8468..f6e5fd15 100644 --- a/dataretrieval/waterdata/stats.py +++ b/dataretrieval/waterdata/stats.py @@ -67,37 +67,11 @@ def _handle_nesting( geopandas branch. Skipping the GeoJSON envelope keeps newly-added fields like ``geometry.type`` from leaking into the result. """ - if body is None: - return _empty_feature_frame(geopd) - - # An empty (or missing) features list — a real mid-pagination - # shape — would otherwise crash the downstream merge with - # ``KeyError: 'monitoring_location_id'`` because neither df nor - # dat would carry the merge key. ``_empty_feature_frame`` bails out - # with a geo-typed empty frame so a later ``pd.concat`` with non-empty - # geo pages doesn't downgrade to a plain DataFrame and strip geometry/CRS. - features = body.get("features") or [] - if not features: + features = _extract_features(body) + if features is None: return _empty_feature_frame(geopd) - # The geopd-missing warning is emitted once at import (see engine module); - # doing it here would log per page. - if not geopd: - outer_props = [ - {k: v for k, v in (f.get("properties") or {}).items() if k != "data"} - for f in features - ] - df = pd.json_normalize(outer_props, sep=".") - df.columns = df.columns.str.split(".").str[-1] - # Stats features don't carry a top-level ``id`` field — the - # geopandas branch (``GeoDataFrame.from_features``) doesn't - # surface one either, so the non-geopd branch stays - # consistent by NOT adding an id column. - _attach_coordinates(df, features) - else: - # Stats features may omit ``geometry`` entirely; ``_geo_feature_frame`` - # is the shared home for that upstream-schema workaround. - df = _geo_feature_frame(features).drop(columns=["data"], errors="ignore") + df = _build_outer_frame(features, geopd) # Unnest json features, properties, data, and values while retaining necessary # metadata to merge with main dataframe. @@ -118,6 +92,44 @@ def _handle_nesting( return df.merge(dat, on="monitoring_location_id", how="left") +def _extract_features(body: dict[str, Any] | None) -> list[dict[str, Any]] | None: + """Return the features list from a response body, or None for empty/missing. + + ``None`` signals the caller to return an empty frame. An empty (or + missing) features list — a real mid-pagination shape — would otherwise + crash the downstream merge with ``KeyError: 'monitoring_location_id'`` + because neither frame would carry the merge key. + """ + if body is None: + return None + return body.get("features") or None + + +def _build_outer_frame(features: list[dict[str, Any]], geopd: bool) -> pd.DataFrame: + """Build the outer (per-monitoring-location) frame from GeoJSON features. + + The geopd-missing warning is emitted once at import (see engine module); + doing it here would log per page. + """ + if not geopd: + outer_props = [ + {k: v for k, v in (f.get("properties") or {}).items() if k != "data"} + for f in features + ] + df = pd.json_normalize(outer_props, sep=".") + df.columns = df.columns.str.split(".").str[-1] + # Stats features don't carry a top-level ``id`` field — the + # geopandas branch (``GeoDataFrame.from_features``) doesn't + # surface one either, so the non-geopd branch stays + # consistent by NOT adding an id column. + _attach_coordinates(df, features) + return df + + # Stats features may omit ``geometry`` entirely; ``_geo_feature_frame`` + # is the shared home for that upstream-schema workaround. + return _geo_feature_frame(features).drop(columns=["data"], errors="ignore") + + def _expand_percentiles(df: pd.DataFrame) -> pd.DataFrame: """Explode percentile value and threshold lists into one row per element. diff --git a/pyproject.toml b/pyproject.toml index 1147e688..54d8ab34 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -47,7 +47,7 @@ include = ["dataretrieval*"] dataretrieval = ["py.typed"] [tool.complexipy] -max-complexity-allowed = 25 +max-complexity-allowed = 10 failed = true [project.optional-dependencies]