From 66ee707dc837b142f4ea2e7e69874c542ff0da51 Mon Sep 17 00:00:00 2001
From: gantian127 <jamy127@foxmail.com>
Date: Tue, 29 Oct 2024 18:19:40 -0600
Subject: [PATCH] update utils function and unit test to support get time info
 from time variable defined as 'date'

---
 src/bmi_era5/utils.py | 54 ++++++++++++++++++++++++++++++++-----------
 tests/util_test.py    | 42 ++++++++++++++++++++++++++++++++-
 2 files changed, 82 insertions(+), 14 deletions(-)

diff --git a/src/bmi_era5/utils.py b/src/bmi_era5/utils.py
index 474f356..055debd 100755
--- a/src/bmi_era5/utils.py
+++ b/src/bmi_era5/utils.py
@@ -1,8 +1,11 @@
 from __future__ import annotations
 
 import os.path
+from datetime import datetime
 
 import cdsapi
+import cftime
+import numpy as np
 import xarray as xr
 
 
@@ -81,19 +84,44 @@ def get_time_info(self):
 
         # time values are float in BMI time function
         if self._data:
-            time_info = {
-                "start_time": float(self._data.valid_time.values[0]),
-                "time_step": 0.0
-                if len(self._data.valid_time.values) == 1
-                else float(
-                    self._data.valid_time.values[1] - self._data.valid_time.values[0]
-                ),
-                "end_time": float(self._data.valid_time.values[-1]),
-                "total_steps": len(self._data.valid_time.values),
-                "time_units": self._data.valid_time.units,
-                "calendar": self._data.valid_time.calendar,
-                "time_value": self._data.valid_time.values.astype("float"),
-            }
+            if "valid_time" in self._data.keys():
+                time_info = {
+                    "start_time": float(self._data.valid_time.values[0]),
+                    "time_step": 0.0
+                    if len(self._data.valid_time.values) == 1
+                    else float(
+                        self._data.valid_time.values[1]
+                        - self._data.valid_time.values[0]
+                    ),
+                    "end_time": float(self._data.valid_time.values[-1]),
+                    "total_steps": len(self._data.valid_time.values),
+                    "time_units": self._data.valid_time.units,
+                    "calendar": self._data.valid_time.calendar,
+                    "time_value": self._data.valid_time.values.astype("float"),
+                }
+            elif "date" in self._data.keys():
+                # convert date time to CF convention values
+                date_objs = [
+                    datetime.strptime(str(date_value), "%Y%m%d")
+                    for date_value in self._data.date.values
+                ]
+                time_units = "seconds since 1970-01-01"
+                calendar = "proleptic_gregorian"
+                cf_dates = cftime.date2num(
+                    date_objs, units=time_units, calendar=calendar
+                )
+
+                time_info = {
+                    "start_time": float(cf_dates[0]),
+                    "time_step": 0.0
+                    if len(cf_dates) == 1
+                    else float(cf_dates[1] - cf_dates[0]),
+                    "end_time": float(cf_dates[-1]),
+                    "total_steps": len(cf_dates),
+                    "time_units": time_units,
+                    "calendar": calendar,
+                    "time_value": np.array(cf_dates, dtype=float),
+                }
 
         return time_info
 
diff --git a/tests/util_test.py b/tests/util_test.py
index d7124ba..c2baaa6 100755
--- a/tests/util_test.py
+++ b/tests/util_test.py
@@ -24,6 +24,23 @@
     )
 ]
 
+parameters2 = [
+    (
+        "reanalysis-era5-single-levels-monthly-means",
+        "monthly_mean.nc",
+        {
+            "product_type": ["monthly_averaged_reanalysis"],
+            "variable": ["2m_dewpoint_temperature"],
+            "year": ["2022"],
+            "month": ["01", "02", "03", "04"],
+            "time": ["00:00"],
+            "data_format": "netcdf",
+            "download_format": "unarchived",
+            "area": [39, -106, 36, -103],
+        },
+    )
+]
+
 
 @pytest.mark.parametrize("name, file, era5_req", parameters)
 def test_get_data(tmpdir, name, file, era5_req):
@@ -76,7 +93,8 @@ def test_get_var_info(tmpdir, name, file, era5_req):
 
 
 @pytest.mark.parametrize("name, file, era5_req", parameters)
-def test_get_time_info(tmpdir, name, file, era5_req):
+def test_get_time_info_valid_time(tmpdir, name, file, era5_req):
+    """Test when time variable is valid_time"""
     path = os.path.join(tmpdir, file)
 
     era5 = Era5Data()
@@ -93,3 +111,25 @@ def test_get_time_info(tmpdir, name, file, era5_req):
     assert time_info_2["total_steps"] == 3
     assert time_info_2["time_units"] == "seconds since 1970-01-01"
     assert time_info_2["calendar"] == "proleptic_gregorian"
+
+
+@pytest.mark.parametrize("name, file, era5_req", parameters2)
+def test_get_time_info_date(tmpdir, name, file, era5_req):
+    """Test when time variable is date"""
+
+    path = os.path.join(tmpdir, file)
+
+    era5 = Era5Data()
+    time_info_1 = era5.get_time_info()
+
+    assert time_info_1 == {}
+
+    era5.get_data(name, era5_req, path)
+    time_info_2 = era5.get_time_info()
+
+    assert time_info_2["start_time"] == 1640995200.0
+    assert time_info_2["end_time"] == 1648771200.0
+    assert time_info_2["time_step"] == 2678400.0
+    assert time_info_2["total_steps"] == 4
+    assert time_info_2["time_units"] == "seconds since 1970-01-01"
+    assert time_info_2["calendar"] == "proleptic_gregorian"