From 75d841b5359c5d79a738016eaad4416af95588fe Mon Sep 17 00:00:00 2001
From: Lukas Pilz <Lukas.Pilz@iup.uni-heidelberg.de>
Date: Thu, 25 Apr 2024 14:24:55 +0200
Subject: [PATCH] Added XTIME fallback for Time coord generation

---
 tests/test_postprocess.py | 18 ++++++++++++++++++
 xwrf/postprocess.py       | 26 ++++++++++++++++++--------
 2 files changed, 36 insertions(+), 8 deletions(-)

diff --git a/tests/test_postprocess.py b/tests/test_postprocess.py
index 4cfb506..d611259 100644
--- a/tests/test_postprocess.py
+++ b/tests/test_postprocess.py
@@ -121,6 +121,24 @@ def test_xtime_handling(sample_dataset_with_kwargs, xtime_dtype):
     assert np.issubdtype(dataset.XTIME.dtype, xtime_dtype)
 
 
+@pytest.mark.parametrize(
+    'sample_dataset_with_kwargs,xtime_dtype',
+    [
+        (('wrfout', {'decode_times': True}), np.datetime64),
+        pytest.param(
+            ('wrfout', {'decode_times': False}),
+            np.datetime64,
+            marks=pytest.mark.xfail(raises=ValueError, strict=True),
+        ),
+    ],
+    indirect=['sample_dataset_with_kwargs'],
+)
+def test_xtime_fallback(sample_dataset_with_kwargs, xtime_dtype):
+    dataset_fallback = xwrf.postprocess._decode_times(sample_dataset_with_kwargs.drop_vars('Times'))
+    dataset = xwrf.postprocess._decode_times(sample_dataset_with_kwargs)
+    assert dataset_fallback.Time.equals(dataset.Time)
+
+
 @pytest.mark.parametrize('sample_dataset', ['lambert_conformal'], indirect=True)
 def test_assign_coord_to_dim_of_different_name(sample_dataset):
     dataset = sample_dataset.pipe(xwrf.postprocess._collapse_time_dim).pipe(
diff --git a/xwrf/postprocess.py b/xwrf/postprocess.py
index 2d7ccf2..e2697b5 100644
--- a/xwrf/postprocess.py
+++ b/xwrf/postprocess.py
@@ -14,14 +14,24 @@ def _decode_times(ds: xr.Dataset) -> xr.Dataset:
     """
     Decode the time variable to datetime64.
     """
-    try:
-        _time = pd.to_datetime(
-            ds.Times.data.astype('str'), errors='raise', format='%Y-%m-%d_%H:%M:%S'
-        )
-    except ValueError:
-        _time = pd.to_datetime(
-            ds.Times.data.astype('str'), errors='raise', format='%Y-%m-%dT%H:%M:%S'
-        )
+    if 'Times' in ds:
+        try:
+            _time = pd.to_datetime(
+                ds.Times.data.astype('str'), errors='raise', format='%Y-%m-%d_%H:%M:%S'
+            )
+        except ValueError:
+            _time = pd.to_datetime(
+                ds.Times.data.astype('str'), errors='raise', format='%Y-%m-%dT%H:%M:%S'
+            )
+    elif 'XTIME' in ds:
+        if ds.XTIME.dtype == 'datetime64[ns]':
+            _time = ds.XTIME.data
+        else:
+            raise ValueError(
+                'XTIME is not in datetime64 format, please use `xr.open_dataset(..., decode_times=True)`.'
+            )
+    else:
+        raise ValueError('No time variable found in the dataset.')
     ds = ds.assign_coords({'Time': _time})
     ds.Time.attrs = {'long_name': 'Time', 'standard_name': 'time'}
     return ds