Skip to content

get_data now returns time indexed dataframes #42

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
kratzert merged 1 commit into from
Oct 15, 2025
Merged

get_data now returns time indexed dataframes #42

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
4 changes: 2 additions & 2 deletions examples/test_australia_fetcher.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -23,9 +23,9 @@
if not data.empty:
print(f"Data for {gauge_id}:")
print(data.head())
print(f"Time series from {data[constants.TIME_INDEX].min()} to {data[constants.TIME_INDEX].max()}")
print(f"Time series from {data.index.min()} to {data.index.max()}")
plt.plot(
data[constants.TIME_INDEX],
data.index,
data[constants.DISCHARGE],
label=gauge_id,
marker="o",
Expand Down
4 changes: 2 additions & 2 deletions examples/test_canada_fetcher.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -20,9 +20,9 @@
if not data.empty:
print(f"Data for {gauge_id}:")
print(data.head())
print(f"Time series from {data[constants.TIME_INDEX].min()} to {data[constants.TIME_INDEX].max()}")
print(f"Time series from {data.index.min()} to {data.index.max()}")
plt.plot(
data[constants.TIME_INDEX],
data.index,
data[constants.DISCHARGE],
label=gauge_id,
marker=".",
Expand Down
4 changes: 2 additions & 2 deletions examples/test_chile_fetcher.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -20,9 +20,9 @@
if not data.empty:
print(f"Data for {gauge_id}:")
print(data.head())
print(f"Time series from {data[constants.TIME_INDEX].min()} to {data[constants.TIME_INDEX].max()}")
print(f"Time series from {data.index.min()} to {data.index.max()}")
plt.plot(
data[constants.TIME_INDEX],
data.index,
data[constants.DISCHARGE],
label=gauge_id,
marker=".",
Expand Down
6 changes: 3 additions & 3 deletions examples/test_france_fetcher.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -20,15 +20,15 @@
if not data.empty:
print(f"Data for {gauge_id}:")
print(data.head())
print(f"Time series from {data[constants.TIME_INDEX].min()} to {data[constants.TIME_INDEX].max()}")
print(f"Time series from {data.index.min()} to {data.index.max()}")
plt.plot(
data[constants.TIME_INDEX],
data.index,
data[constants.DISCHARGE],
label=gauge_id,
marker=".",
linestyle="-",
)
plt.xlim(data[constants.TIME_INDEX].min(), data[constants.TIME_INDEX].max())
plt.xlim(data.index.min(), data.index.max())
else:
print(f"No data found for {gauge_id}")

Expand Down
4 changes: 2 additions & 2 deletions examples/test_japan_fetcher.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -22,9 +22,9 @@
if not data.empty:
print(f"Data for {gauge_id}:")
print(data.head())
print(f"Time series from {data[constants.TIME_INDEX].min()} to {data[constants.TIME_INDEX].max()}")
print(f"Time series from {data.index.min()} to {data.index.max()}")
plt.plot(
data[constants.TIME_INDEX],
data.index,
data[constants.DISCHARGE],
label=gauge_id,
marker="o",
Expand Down
4 changes: 2 additions & 2 deletions examples/test_poland_fetcher.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -29,9 +29,9 @@
if not data.empty:
print(f"Data for {gauge_id}:")
print(data.head())
print(f"Time series from {data[constants.TIME_INDEX].min()} to {data[constants.TIME_INDEX].max()}")
print(f"Time series from {data.index.min()} to {data.index.max()}")
plt.plot(
data[constants.TIME_INDEX],
data.index,
data[variable],
label=gauge_id,
marker=".",
Expand Down
4 changes: 2 additions & 2 deletions examples/test_slovenia_fetcher.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -20,9 +20,9 @@
if not data.empty:
print(f"Data for {gauge_id}:")
print(data.head())
print(f"Time series from {data[constants.TIME_INDEX].min()} to {data[constants.TIME_INDEX].max()}")
print(f"Time series from {data.index.min()} to {data.index.max()}")
plt.plot(
data[constants.TIME_INDEX],
data.index,
data[constants.DISCHARGE],
label=gauge_id,
marker="o",
Expand Down
4 changes: 2 additions & 2 deletions examples/test_southafrica_fetcher.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -20,9 +20,9 @@
if not data.empty:
print(f"Data for {gauge_id}:")
print(data.head())
print(f"Time series from {data[constants.TIME_INDEX].min()} to {data[constants.TIME_INDEX].max()}")
print(f"Time series from {data.index.min()} to {data.index.max()}")
plt.plot(
data[constants.TIME_INDEX],
data.index,
data[constants.DISCHARGE],
label=gauge_id,
marker=".",
Expand Down
2 changes: 1 addition & 1 deletion examples/test_uk_fetcher.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -19,7 +19,7 @@
print(f"Data for {gauge_id}:")
print(data.head())
plt.plot(
data[constants.TIME_INDEX],
data.index,
data[constants.DISCHARGE],
label=gauge_id.split("/")[-1],
)
Expand Down
4 changes: 2 additions & 2 deletions examples/test_uk_nrfa_fetcher.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -31,9 +31,9 @@
if not data.empty:
print(f"Data for {gauge_id}:")
print(data.head())
print(f"Time series from {data[constants.TIME_INDEX].min()} to {data[constants.TIME_INDEX].max()}")
print(f"Time series from {data.index.min()} to {data.index.max()}")
plt.plot(
data[constants.TIME_INDEX],
data.index,
data[constants.DISCHARGE],
label=gauge_id,
marker=".",
Expand Down
4 changes: 2 additions & 2 deletions examples/test_usa_fetcher.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -23,9 +23,9 @@
if not data.empty:
print(f"Data for {gauge_id}:")
print(data.head())
print(f"Time series from {data[constants.TIME_INDEX].min()} to {data[constants.TIME_INDEX].max()}")
print(f"Time series from {data.index.min()} to {data.index.max()}")
plt.plot(
data[constants.TIME_INDEX],
data.index,
data[constants.DISCHARGE],
label=gauge_id,
marker="o",
Expand Down
2 changes: 1 addition & 1 deletion rivretrieve/australia.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -148,7 +148,7 @@ def _parse_data(self, gauge_id: str, raw_data: Optional[str], variable: str) ->
df["Value"] = pd.to_numeric(df["Value"], errors="coerce")
df = df.rename(columns={"Value": variable})
df[constants.TIME_INDEX] = pd.to_datetime(df[constants.TIME_INDEX])
return df[[constants.TIME_INDEX, variable]].dropna()
return df[[constants.TIME_INDEX, variable]].dropna().set_index(constants.TIME_INDEX)
except Exception as e:
logger.error(f"Error parsing CSV data for site {gauge_id}: {e}")
return pd.DataFrame(columns=[constants.TIME_INDEX, variable])
Expand Down
6 changes: 4 additions & 2 deletions rivretrieve/base.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -25,12 +25,14 @@ def get_data(

Args:
gauge_id: The site-specific identifier for the gauge.
variable: The variable to fetch (e.g., 'stage' or 'discharge').
variable: The variable to fetch, should be one of the values from constants.py
(e.g., constants.DISCHARGE, constants.STAGE).
start_date: Optional start date in 'YYYY-MM-DD' format.
end_date: Optional end date in 'YYYY-MM-DD' format.

Returns:
A pandas DataFrame with 'Date' and the variable column ('H' or 'Q').
A pandas DataFrame indexed by time (constants.TIME_INDEX) with a column
for the requested variable (e.g., constants.DISCHARGE).
"""
pass

Expand Down
2 changes: 1 addition & 1 deletion rivretrieve/canada.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -186,7 +186,7 @@ def get_data(
df_long[[constants.TIME_INDEX, variable]]
.dropna()
.sort_values(by=constants.TIME_INDEX)
.reset_index(drop=True)
.set_index(constants.TIME_INDEX)
)

except Exception as e:
Expand Down
4 changes: 2 additions & 2 deletions rivretrieve/chile.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -104,7 +104,7 @@ def _parse_data(
df[[constants.TIME_INDEX, variable]]
.dropna()
.sort_values(by=constants.TIME_INDEX)
.reset_index(drop=True)
.set_index(constants.TIME_INDEX)
)
except Exception as e:
logger.error(f"Error parsing data for site {gauge_id}: {e}")
Expand Down Expand Up @@ -135,7 +135,7 @@ def get_data(
# Filter by date range
start_date_dt = pd.to_datetime(start_date)
end_date_dt = pd.to_datetime(end_date)
df = df[(df[constants.TIME_INDEX] >= start_date_dt) & (df[constants.TIME_INDEX] <= end_date_dt)]
df = df[(df.index >= start_date_dt) & (df.index <= end_date_dt)]
return df
except Exception as e:
logger.error(f"Failed to get data for site {gauge_id}, variable {variable}: {e}")
Expand Down
2 changes: 1 addition & 1 deletion rivretrieve/france.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -107,7 +107,7 @@ def _parse_data(
df[[constants.TIME_INDEX, variable]]
.dropna()
.sort_values(by=constants.TIME_INDEX)
.reset_index(drop=True)
.set_index(constants.TIME_INDEX)
)
except Exception as e:
logger.error(f"Error parsing JSON data for site {gauge_id}: {e}")
Expand Down
6 changes: 3 additions & 3 deletions rivretrieve/japan.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -141,9 +141,9 @@ def _parse_data(

final_df = pd.concat(all_dfs, ignore_index=True)
final_df = final_df.rename(columns={"Value": variable})
final_df = final_df.sort_values(by=constants.TIME_INDEX).reset_index(drop=True)
final_df = final_df.sort_values(by=constants.TIME_INDEX)

return final_df
return final_df.set_index(constants.TIME_INDEX)

def get_data(
self,
Expand All @@ -164,7 +164,7 @@ def get_data(

start_date_dt = pd.to_datetime(start_date)
end_date_dt = pd.to_datetime(end_date)
df = df[(df[constants.TIME_INDEX] >= start_date_dt) & (df[constants.TIME_INDEX] <= end_date_dt)]
df = df[(df.index >= start_date_dt) & (df.index <= end_date_dt)]
return df

except Exception as e:
Expand Down
2 changes: 1 addition & 1 deletion rivretrieve/poland.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -209,7 +209,7 @@ def get_data(

data_array = ds[variable].sel(gauge_id=gauge_id, time=slice(start_date, end_date))
df = data_array.to_pandas().dropna().reset_index().rename(columns={variable: variable})
return df[[constants.TIME_INDEX, variable]]
return df.set_index(constants.TIME_INDEX)[[variable]]

except KeyError:
logger.info(f"No data found for gauge {gauge_id} in the selected date range.")
Expand Down
4 changes: 2 additions & 2 deletions rivretrieve/slovenia.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -85,7 +85,7 @@ def _parse_data(self, gauge_id: str, raw_data: Optional[str], variable: str) ->
df[[constants.TIME_INDEX, variable]]
.dropna()
.sort_values(by=constants.TIME_INDEX)
.reset_index(drop=True)
.set_index(constants.TIME_INDEX)
)

except Exception as e:
Expand Down Expand Up @@ -113,7 +113,7 @@ def get_data(
# Filter by date range
start_date_dt = pd.to_datetime(start_date)
end_date_dt = pd.to_datetime(end_date)
df = df[(df[constants.TIME_INDEX] >= start_date_dt) & (df[constants.TIME_INDEX] <= end_date_dt)]
df = df[(df.index >= start_date_dt) & (df.index <= end_date_dt)]
return df
except Exception as e:
logger.error(f"Failed to get data for site {gauge_id}, variable {variable}: {e}")
Expand Down
2 changes: 1 addition & 1 deletion rivretrieve/southafrica.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -155,7 +155,7 @@ def _parse_data(
daily_df = full_df[[constants.TIME_INDEX, "D_AVG_FR"]].rename(columns={"D_AVG_FR": "Value"})

daily_df = daily_df.rename(columns={"Value": variable})
return daily_df.dropna().sort_values(by=constants.TIME_INDEX).reset_index(drop=True)
return daily_df.dropna().sort_values(by=constants.TIME_INDEX).set_index(constants.TIME_INDEX)

except Exception as e:
logger.error(f"Error parsing data for site {gauge_id}: {e}")
Expand Down
4 changes: 2 additions & 2 deletions rivretrieve/uk.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -126,7 +126,7 @@ def _parse_data(self, gauge_id: str, raw_data: List[Dict[str, Any]], variable: s
complete_ts = pd.DataFrame(date_range, columns=[constants.TIME_INDEX])
df_daily = pd.merge(complete_ts, df_daily, on=constants.TIME_INDEX, how="left")

return df_daily
return df_daily.set_index(constants.TIME_INDEX)

def get_data(
self,
Expand All @@ -148,7 +148,7 @@ def get_data(
# Filter by exact start and end date after processing
start_date_dt = pd.to_datetime(start_date)
end_date_dt = pd.to_datetime(end_date)
df = df[(df[constants.TIME_INDEX] >= start_date_dt) & (df[constants.TIME_INDEX] <= end_date_dt)]
df = df[(df.index >= start_date_dt) & (df.index <= end_date_dt)]
return df

except Exception as e:
Expand Down
4 changes: 2 additions & 2 deletions rivretrieve/uk_nrfa.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -102,7 +102,7 @@ def _parse_data(self, gauge_id: str, raw_data: Optional[Dict[str, Any]], variabl
df[constants.TIME_INDEX] = pd.to_datetime(df["time"], format="ISO8601").dt.date
df[constants.TIME_INDEX] = pd.to_datetime(df[constants.TIME_INDEX])
df[variable] = pd.to_numeric(df[variable], errors="coerce")
return df[[constants.TIME_INDEX, variable]].dropna().reset_index(drop=True)
return df[[constants.TIME_INDEX, variable]].dropna().set_index(constants.TIME_INDEX)
except Exception as e:
logger.error(f"Error parsing NRFA data for {gauge_id}: {e}")
return pd.DataFrame(columns=[constants.TIME_INDEX, variable])
Expand All @@ -128,7 +128,7 @@ def get_data(
# Filter by date range
start_date_dt = pd.to_datetime(start_date)
end_date_dt = pd.to_datetime(end_date)
df = df[(df[constants.TIME_INDEX] >= start_date_dt) & (df[constants.TIME_INDEX] <= end_date_dt)]
df = df[(df.index >= start_date_dt) & (df.index <= end_date_dt)]
return df
except Exception as e:
logger.error(f"Failed to get data for site {gauge_id}, variable {variable}: {e}")
Expand Down
2 changes: 1 addition & 1 deletion rivretrieve/usa.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -76,7 +76,7 @@ def _parse_data(self, gauge_id: str, raw_data: pd.DataFrame, variable: str) -> p
mult = 0.0283168466
df[variable] = pd.to_numeric(df[value_col], errors="coerce") * mult

return df[[constants.TIME_INDEX, variable]].dropna()
return df[[constants.TIME_INDEX, variable]].dropna().set_index(constants.TIME_INDEX)

def get_data(
self,
Expand Down
2 changes: 1 addition & 1 deletion tests/test_australia.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -47,7 +47,7 @@ def bom_request_side_effect(params):
constants.TIME_INDEX: pd.to_datetime(["2010-01-01", "2010-01-02", "2010-01-03"]),
constants.DISCHARGE: [0.000, 3.710, 3.211],
}
expected_df = pd.DataFrame(expected_data)
expected_df = pd.DataFrame(expected_data).set_index(constants.TIME_INDEX)

assert_frame_equal(result_df, expected_df)
self.assertEqual(mock_make_bom_request.call_count, 2)
Expand Down
4 changes: 2 additions & 2 deletions tests/test_canada.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -36,7 +36,7 @@ def test_get_data_discharge(self, mock_hydat_path, mock_download, mock_requests)
),
constants.DISCHARGE: [1.1, 1.2, 1.3, 1.4, 1.5],
}
expected_df = pd.DataFrame(expected_data)
expected_df = pd.DataFrame(expected_data).set_index(constants.TIME_INDEX)

assert_frame_equal(result_df, expected_df)

Expand All @@ -62,7 +62,7 @@ def test_get_data_stage(self, mock_hydat_path, mock_download, mock_requests):
),
constants.STAGE: [10.1, 10.2, 10.3, 10.4, 10.5],
}
expected_df = pd.DataFrame(expected_data)
expected_df = pd.DataFrame(expected_data).set_index(constants.TIME_INDEX)

assert_frame_equal(result_df, expected_df)

Expand Down
2 changes: 1 addition & 1 deletion tests/test_chile.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -50,7 +50,7 @@ def get_side_effect(*args, **kwargs):
constants.TIME_INDEX: pd.to_datetime(["2022-01-01", "2022-01-02", "2022-01-03"]),
constants.DISCHARGE: [15.5, 16.0, 15.8],
}
expected_df = pd.DataFrame(expected_data)
expected_df = pd.DataFrame(expected_data).set_index(constants.TIME_INDEX)

assert_frame_equal(result_df, expected_df)
self.assertEqual(mock_get.call_count, 2)
Expand Down
2 changes: 1 addition & 1 deletion tests/test_france.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -38,7 +38,7 @@ def test_get_data_discharge(self, mock_get):
constants.TIME_INDEX: pd.to_datetime(["2023-01-01", "2023-01-02", "2023-01-03"]),
constants.DISCHARGE: [15.0005, 16.0000, 15.5002], # Divided by 1000
}
expected_df = pd.DataFrame(expected_data)
expected_df = pd.DataFrame(expected_data).set_index(constants.TIME_INDEX)

assert_frame_equal(result_df, expected_df)
mock_get.assert_called_once()
Expand Down
4 changes: 2 additions & 2 deletions tests/test_japan.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -40,9 +40,9 @@ def test_get_data_discharge(self, mock_get):
constants.TIME_INDEX: expected_dates,
constants.DISCHARGE: expected_values,
}
expected_df = pd.DataFrame(expected_data)
expected_df = pd.DataFrame(expected_data).set_index(constants.TIME_INDEX)

assert_frame_equal(result_df.reset_index(drop=True), expected_df)
assert_frame_equal(result_df, expected_df)
mock_get.assert_called_once()
# Check that the params are correct
mock_args, mock_kwargs = mock_get.call_args
Expand Down
8 changes: 4 additions & 4 deletions tests/test_poland.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -46,9 +46,9 @@ def test_get_data_discharge(self, mock_create_cache):
constants.TIME_INDEX: expected_dates,
constants.DISCHARGE: expected_values,
}
expected_df = pd.DataFrame(expected_data)
expected_df = pd.DataFrame(expected_data).set_index(constants.TIME_INDEX)

assert_frame_equal(result_df.reset_index(drop=True), expected_df)
assert_frame_equal(result_df, expected_df)
mock_create_cache.assert_not_called() # Cache should not be recreated

@patch("rivretrieve.poland.PolandFetcher._create_cache") # Mock cache creation
Expand All @@ -67,9 +67,9 @@ def test_get_data_stage(self, mock_create_cache):
constants.TIME_INDEX: expected_dates,
constants.STAGE: expected_values,
}
expected_df = pd.DataFrame(expected_data)
expected_df = pd.DataFrame(expected_data).set_index(constants.TIME_INDEX)

assert_frame_equal(result_df.reset_index(drop=True), expected_df)
assert_frame_equal(result_df, expected_df)
mock_create_cache.assert_not_called()

@patch("rivretrieve.utils.requests_retry_session")
Expand Down
Loading