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