feat: add function to save and load locations to csv file

cwmscli/load/location/location.py

@@ -5,6 +5,7 @@
 
 from cwmscli import requirements as reqs
 from cwmscli.load.root import (
+    csv_source_target_options,
     load_group,
     shared_source_target_options,
     validate_cda_targets,
@@ -24,12 +25,14 @@ def location(ctx):
 @location.command(
     "ids-all",
     help=(
-        "Copy locations from a source CDA catalog to a target CDA. "
+        "Copy locations from a source CDA catalog to a target CDA, "
+        "or to/from a CSV file via --source-csv / --target-csv. "
         "The --like and --location-kind-like filters use CDA regex semantics. "
         f"Regex guide: {CDA_REGEXP_GUIDE_URL}"
     ),
 )
 @shared_source_target_options
+@csv_source_target_options(allow_source_csv=True, allow_target_csv=True)
 @click.option(
     "--like",
     default=None,
@@ -57,14 +60,16 @@ def location(ctx):
 @requires(reqs.cwms)
 @validate_cda_targets
 def load_locations(
-    source_cda: str,
-    source_office: str,
-    target_cda: str,
+    source_cda: Optional[str],
+    source_office: Optional[str],
+    target_cda: Optional[str],
     target_api_key: Optional[str],
     verbose: int,
     dry_run: bool,
     like: Optional[str],
     location_kind_like: Optional[Iterable[str]] = None,
+    source_csv: Optional[str] = None,
+    target_csv: Optional[str] = None,
 ):
     from cwmscli.load.location.location_ids import load_locations as _load_locations
 
@@ -77,14 +82,20 @@ def load_locations(
         dry_run=dry_run,
         like=like,
         location_kind_like=location_kind_like,
+        source_csv=source_csv,
+        target_csv=target_csv,
     )
 
 
 @location.command(
     "ids-bygroup",
-    help="Copy locations from a CWMS Location Group (source CDA) to a target CDA.",
+    help=(
+        "Copy locations from a CWMS Location Group (source CDA) to a target CDA, "
+        "or export the resolved members to a CSV file via --target-csv."
+    ),
 )
 @shared_source_target_options
+@csv_source_target_options(allow_source_csv=False, allow_target_csv=True)
 @click.option(
     "--group-id", required=True, help="Location Group ID (e.g., 'Ark Basin')."
 )
@@ -112,7 +123,7 @@ def load_locations(
 def load_locations_from_group(
     source_cda: str,
     source_office: str,
-    target_cda: str,
+    target_cda: Optional[str],
     target_api_key: Optional[str],
     verbose: int,
     group_id: str,
@@ -121,6 +132,7 @@ def load_locations_from_group(
     category_office_id: Optional[str],
     filter_office: bool,
     dry_run: bool,
+    target_csv: Optional[str] = None,
 ):
     from cwmscli.load.location.location_ids_bygroup import copy_from_group
 
@@ -136,4 +148,5 @@ def load_locations_from_group(
         category_office_id=category_office_id,
         filter_office=filter_office,
         dry_run=dry_run,
+        target_csv=target_csv,
     )

cwmscli/load/location/location_ids.py

@@ -1,9 +1,11 @@
 import logging
+import math
 import re
 from typing import Iterable, Optional
 
 import click
 import cwms
+import pandas as pd
 
 from cwmscli.utils import init_cwms_session
 from cwmscli.utils.links import CDA_REGEXP_GUIDE_URL
@@ -12,18 +14,23 @@
 
 
 def load_locations(
-    source_cda: str,
-    source_office: str,
-    target_cda: str,
+    source_cda: Optional[str],
+    source_office: Optional[str],
+    target_cda: Optional[str],
     target_api_key: Optional[str],
     verbose: int,
     dry_run: bool,
     like: Optional[str],
     location_kind_like: Optional[Iterable[str]] = "ALL",
+    source_csv: Optional[str] = None,
+    target_csv: Optional[str] = None,
 ):
+    src_label = source_csv or source_cda or "-"
+    tgt_label = target_csv or target_cda or "-"
+
     if verbose:
         logger.info(
-            f"[load locations] source={source_cda} ({source_office}) -> target={target_cda}"
+            f"[load locations] source={src_label} ({source_office or '-'}) -> target={tgt_label}"
         )
         logger.info(
             f"  like={like or '-'}  kinds={list(location_kind_like) or '-'}  dry_run={dry_run}"
@@ -35,63 +42,41 @@ def load_locations(
         ):
             logger.info("  CDA regex guide: %s", CDA_REGEXP_GUIDE_URL)
 
-    init_cwms_session(cwms, api_root=source_cda)
-
-    cat_kwargs = {"office_id": source_office}
-    if like:
-        cat_kwargs["like"] = like
-    kinds = list(location_kind_like) if location_kind_like else ["ALL"]
-    if "ALL" in kinds:
-        kinds = ["ALL"]
-
-    locations = []
-
-    if kinds == ["ALL"] and not like:
-        locations = cwms.get_locations(office_id=source_office).json
+    if source_csv:
+        df = pd.read_csv(source_csv)
+        locations = df.to_dict(orient="records")
     else:
-        seen_location_ids = set()
-        for kind in kinds:
-            cat_kwargs_k = dict(cat_kwargs)
-            if kind != "ALL":
-                cat_kwargs_k["location_kind_like"] = kind
-
-            if verbose >= 2:
-                logger.debug("  > catalog query: %s", cat_kwargs_k)
-
-            resp = cwms.get_locations_catalog(**cat_kwargs_k)
-            if resp.df.empty:
-                continue
-
-            for location_id in resp.df["name"].tolist():
-                if location_id in seen_location_ids:
-                    continue
-                seen_location_ids.add(location_id)
-                if verbose >= 2:
-                    logger.debug("  > location fetch: %s", location_id)
-                detail_resp = cwms.get_locations(
-                    office_id=source_office,
-                    location_ids=rf"^{re.escape(location_id)}$",
-                )
-                if detail_resp and detail_resp.json:
-                    locations.extend(detail_resp.json)
+        init_cwms_session(cwms, api_root=source_cda)
+        locations = _fetch_locations_from_cda(
+            source_office=source_office,
+            like=like,
+            location_kind_like=location_kind_like,
+            verbose=verbose,
+        )
 
     if verbose:
-        logger.info("Fetched %s locations from source", len(locations))
+        logger.info("Got %s locations from source", len(locations))
 
     if dry_run:
         for loc in locations:
             logger.info(
-                f"[dry-run] would store Location(name={loc['name']}) to {target_cda} ({source_office})"
+                f"[dry-run] would store Location(name={loc['name']}) to {tgt_label} "
+                f"({source_office or loc.get('office-id') or '-'})"
             )
         return
 
-    # init target once
+    if target_csv:
+        pd.DataFrame(locations).to_csv(target_csv, index=False)
+        click.echo(f"Wrote {len(locations)} locations to {target_csv}")
+        return
+
     init_cwms_session(cwms, api_root=target_cda, api_key=target_api_key)
 
     errors = 0
     for loc in locations:
+        loc = _clean_row(loc)
         try:
-            if loc["active"] is True:
+            if loc.get("active") is True:
                 result = cwms.store_location(data=loc, fail_if_exists=False)
                 if verbose:
                     logger.info("%s", result)
@@ -103,3 +88,60 @@ def load_locations(
         raise click.ClickException(f"Completed with {errors} error(s).")
 
     click.echo("Done.")
+
+
+def _fetch_locations_from_cda(
+    source_office: str,
+    like: Optional[str],
+    location_kind_like: Optional[Iterable[str]],
+    verbose: int,
+) -> list:
+    cat_kwargs = {"office_id": source_office}
+    if like:
+        cat_kwargs["like"] = like
+    kinds = list(location_kind_like) if location_kind_like else ["ALL"]
+    if "ALL" in kinds:
+        kinds = ["ALL"]
+
+    if kinds == ["ALL"] and not like:
+        return cwms.get_locations(office_id=source_office).json
+
+    locations = []
+    seen_location_ids = set()
+    for kind in kinds:
+        cat_kwargs_k = dict(cat_kwargs)
+        if kind != "ALL":
+            cat_kwargs_k["location_kind_like"] = kind
+
+        if verbose >= 2:
+            logger.debug("  > catalog query: %s", cat_kwargs_k)
+
+        resp = cwms.get_locations_catalog(**cat_kwargs_k)
+        if resp.df.empty:
+            continue
+
+        for location_id in resp.df["name"].tolist():
+            if location_id in seen_location_ids:
+                continue
+            seen_location_ids.add(location_id)
+            if verbose >= 2:
+                logger.debug("  > location fetch: %s", location_id)
+            detail_resp = cwms.get_locations(
+                office_id=source_office,
+                location_ids=rf"^{re.escape(location_id)}$",
+            )
+            if detail_resp and detail_resp.json:
+                locations.extend(detail_resp.json)
+    return locations
+
+
+def _clean_row(loc: dict) -> dict:
+    cleaned = {}
+    for k, v in loc.items():
+        if isinstance(v, float) and math.isnan(v):
+            cleaned[k] = None
+        elif isinstance(v, str) and v in ("True", "False"):
+            cleaned[k] = v == "True"
+        else:
+            cleaned[k] = v
+    return cleaned

cwmscli/load/location/location_ids_bygroup.py

@@ -5,6 +5,7 @@
 
 import click
 import cwms
+import pandas as pd
 
 from cwmscli.utils import init_cwms_session
 
@@ -22,7 +23,7 @@ def exact_or_regex(ids: list[str]) -> str:
 def copy_from_group(
     source_cda: str,
     source_office: str,
-    target_cda: str,
+    target_cda: Optional[str],
     target_api_key: Optional[str],
     verbose: int,
     group_id: str,
@@ -31,6 +32,7 @@ def copy_from_group(
     category_office_id: Optional[str],
     filter_office: bool,
     dry_run: bool,
+    target_csv: Optional[str] = None,
 ):
     group_office_id = group_office_id or source_office
     category_office_id = category_office_id or source_office
@@ -103,10 +105,16 @@ def copy_from_group(
     if dry_run:
         for loc in locations:
             logger.info(
-                f"[dry-run] would store Location(name={loc['name']}) to {target_cda} ({source_office})"
+                f"[dry-run] would store Location(name={loc['name']}) to "
+                f"{target_csv or target_cda} ({source_office})"
             )
         return
 
+    if target_csv:
+        pd.DataFrame(locations).to_csv(target_csv, index=False)
+        click.echo(f"Wrote {len(locations)} locations to {target_csv}")
+        return
+
     try:
         init_cwms_session(cwms, api_root=target_cda, api_key=target_api_key)
     except Exception as e: