Skip to content

Commit cff06d4

Browse files
committed
feat: implement as discussed in PR #8322
1 parent df0ddaf commit cff06d4

File tree

1 file changed

+26
-14
lines changed

1 file changed

+26
-14
lines changed

xarray/coding/times.py

Lines changed: 26 additions & 14 deletions
Original file line numberDiff line numberDiff line change
@@ -628,25 +628,37 @@ def _encode_datetime_with_cftime(dates, units: str, calendar: str) -> np.ndarray
628628
if cftime is None:
629629
raise ModuleNotFoundError("No module named 'cftime'")
630630

631+
dates = np.array(dates)
632+
633+
if dates.shape == ():
634+
dates = dates.reshape(1)
635+
631636
if np.issubdtype(dates.dtype, np.datetime64):
632637
# numpy's broken datetime conversion only works for us precision
633638
dates = dates.astype("M8[us]").astype(datetime)
634639

635-
def encode_datetime(d):
636-
# Since netCDF files do not support storing float128 values, we ensure
637-
# that float64 values are used by setting longdouble=False in num2date.
638-
# This try except logic can be removed when xarray's minimum version of
639-
# cftime is at least 1.6.2.
640-
try:
641-
return (
642-
np.nan
643-
if d is None
644-
else cftime.date2num(d, units, calendar, longdouble=False)
645-
)
646-
except TypeError:
647-
return np.nan if d is None else cftime.date2num(d, units, calendar)
640+
# Find all the None position
641+
none_position = np.equal(dates, None)
642+
643+
# Remove None from the dates and return new array
644+
filtered_dates = dates[~none_position]
648645

649-
return np.array([encode_datetime(d) for d in dates.ravel()]).reshape(dates.shape)
646+
# Since netCDF files do not support storing float128 values, we ensure
647+
# that float64 values are used by setting longdouble=False in num2date.
648+
# This try except logic can be removed when xarray's minimum version of
649+
# cftime is at least 1.6.2.
650+
try:
651+
encoded_nums = cftime.date2num(
652+
filtered_dates, units, calendar, longdouble=False
653+
)
654+
except TypeError:
655+
encoded_nums = cftime.date2num(filtered_dates, units, calendar)
656+
657+
# Create a full matrix of NaN
658+
# And fill the num dates in the not NaN or None position
659+
result = np.full(dates.shape, np.nan)
660+
result[np.nonzero(~none_position)] = encoded_nums
661+
return result
650662

651663

652664
def cast_to_int_if_safe(num) -> np.ndarray:

0 commit comments

Comments
 (0)