add back in geometry when skip_geometry is false but properties do not contain geometry

Author: ehinman

dataretrieval/waterdata/utils.py

@@ -653,7 +653,11 @@ def _arrange_cols(
     df = df.rename(columns={"id": output_id})
 
     # If properties are provided, filter to only those columns
+    # plus geometry if skip_geometry is False
     if properties and not all(pd.isna(properties)):
+        # Make sure geometry stays in the dataframe if skip_geometry is False
+        if 'geometry' in df.columns and 'geometry' not in properties:
+            properties.append('geometry')
         # id is technically a valid column from the service, but these
         # functions make the name more specific. So, if someone requests
         # 'id', give them the output_id column

tests/waterdata_test.py

@@ -214,6 +214,15 @@ def test_get_latest_daily():
     assert hasattr(md, 'url')
     assert hasattr(md, 'query_time')
 
+def test_get_latest_daily_properties_geometry():
+    df, md = get_latest_daily(
+        monitoring_location_id=["USGS-05427718", "USGS-05427719"],
+        parameter_code=["00060", "00065"],
+        properties=['monitoring_location_id', 'parameter_code', 'time', 'value', 'unit_of_measure']
+    )
+    assert "geometry" in df.columns
+    assert df.shape[1] == 6
+
 def test_get_field_measurements():
     df, md = get_field_measurements(
         monitoring_location_id="USGS-05427718",