Skip to content

Commit 3512bee

Browse files
authored
Merge dc2609c into 91a52dc
2 parents 91a52dc + dc2609c commit 3512bee

File tree

1 file changed

+27
-14
lines changed

1 file changed

+27
-14
lines changed

xarray/coding/times.py

Lines changed: 27 additions & 14 deletions
Original file line numberDiff line numberDiff line change
@@ -628,25 +628,38 @@ 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+
original_shape = dates.shape
633+
631634
if np.issubdtype(dates.dtype, np.datetime64):
632635
# numpy's broken datetime conversion only works for us precision
633636
dates = dates.astype("M8[us]").astype(datetime)
634637

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)
638+
dates = np.atleast_1d(dates)
639+
640+
# Find all the None position
641+
none_position = np.equal(dates, None)
642+
filtered_dates = dates[~none_position]
643+
644+
# Since netCDF files do not support storing float128 values, we ensure
645+
# that float64 values are used by setting longdouble=False in num2date.
646+
# This try except logic can be removed when xarray's minimum version of
647+
# cftime is at least 1.6.2.
648+
try:
649+
encoded_nums = cftime.date2num(
650+
filtered_dates, units, calendar, longdouble=False
651+
)
652+
except TypeError:
653+
encoded_nums = cftime.date2num(filtered_dates, units, calendar)
654+
655+
if filtered_dates.size == none_position.size:
656+
return encoded_nums.reshape(original_shape)
648657

649-
return reshape(np.array([encode_datetime(d) for d in ravel(dates)]), dates.shape)
658+
# Create a full matrix of NaN
659+
# And fill the num dates in the not NaN or None position
660+
result = np.full(dates.shape, np.nan)
661+
result[np.nonzero(~none_position)] = encoded_nums
662+
return result.reshape(original_shape)
650663

651664

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

0 commit comments

Comments
 (0)