Update 'stations' API for coops data

pmav99 · Dec 26, 2023 · 13ba6be · 13ba6be
1 parent b60eb98
commit 13ba6be
Show file tree

Hide file tree

Showing 2 changed files with 4 additions and 8 deletions.
diff --git a/searvey/coops.py b/searvey/coops.py
@@ -981,7 +981,7 @@ def get_coops_stations(
     lon_max: Optional[float] = None,
     lat_min: Optional[float] = None,
     lat_max: Optional[float] = None,
-    metadata_source: StationMetadataSource = "nws",
+    metadata_source: Union[StationMetadataSource, str] = "nws",
 ) -> geopandas.GeoDataFrame:
     md_src = StationMetadataSource(metadata_source)
 

diff --git a/searvey/stations.py b/searvey/stations.py
@@ -4,7 +4,6 @@
 from enum import Enum
 
 import geopandas as gpd
-import numpy as np
 import pandas as pd
 from shapely.geometry import MultiPolygon
 from shapely.geometry import Polygon
@@ -92,20 +91,17 @@ def _get_ioc_stations(
 def _get_coops_stations(
     region: Polygon | MultiPolygon | None = None,
 ) -> gpd.GeoDataFrame:
-    coops_gdf = coops.coops_stations_within_region(region=region)
+    coops_gdf = coops.get_coops_stations(region=region, metadata_source="main")
     coops_gdf = coops_gdf.assign(
         provider=Provider.COOPS.value,
         provider_id=coops_gdf.index,
-        country=np.where(coops_gdf.state.str.len() > 0, "USA", None),  # type: ignore[call-overload]
+        country=coops_gdf.state.where(coops_gdf.state.str.len() != 2, "USA"),
         location=coops_gdf[["name", "state"]].agg(", ".join, axis=1),
         lon=coops_gdf.geometry.x,
         lat=coops_gdf.geometry.y,
         is_active=coops_gdf.status == "active",
         start_date=pd.NaT,
-        last_observation=coops_gdf[coops_gdf.status == "discontinued"]
-        .removed.str[:18]
-        .apply(pd.to_datetime)
-        .dt.tz_localize("UTC"),
+        last_observation=coops_gdf[coops_gdf.status == "discontinued"].removed.dt.tz_localize("UTC"),
     )[STATIONS_COLUMNS]
     return coops_gdf