Skip to content

Commit 3f6bf49

Browse files
committed
feat: implement as discussed in PR #8322
1 parent d8b7644 commit 3f6bf49

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
@@ -625,25 +625,37 @@ def _encode_datetime_with_cftime(dates, units: str, calendar: str) -> np.ndarray
625625
if cftime is None:
626626
raise ModuleNotFoundError("No module named 'cftime'")
627627

628+
dates = np.array(dates)
629+
630+
if dates.shape == ():
631+
dates = dates.reshape(1)
632+
628633
if np.issubdtype(dates.dtype, np.datetime64):
629634
# numpy's broken datetime conversion only works for us precision
630635
dates = dates.astype("M8[us]").astype(datetime)
631636

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

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

648660

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

0 commit comments

Comments
 (0)