Skip to content

Commit f28fc7d

Browse files
committed
fix: Stick closer to the original function and add PR suggestions
1 parent 3f6bf49 commit f28fc7d

File tree

1 file changed

+7
-6
lines changed

1 file changed

+7
-6
lines changed

xarray/coding/times.py

Lines changed: 7 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -626,18 +626,16 @@ def _encode_datetime_with_cftime(dates, units: str, calendar: str) -> np.ndarray
626626
raise ModuleNotFoundError("No module named 'cftime'")
627627

628628
dates = np.array(dates)
629-
630-
if dates.shape == ():
631-
dates = dates.reshape(1)
629+
original_shape = dates.shape
632630

633631
if np.issubdtype(dates.dtype, np.datetime64):
634632
# numpy's broken datetime conversion only works for us precision
635633
dates = dates.astype("M8[us]").astype(datetime)
636634

635+
dates = np.atleast_1d(dates)
636+
637637
# Find all the None position
638638
none_position = np.equal(dates, None)
639-
640-
# Remove None from the dates and return new array
641639
filtered_dates = dates[~none_position]
642640

643641
# Since netCDF files do not support storing float128 values, we ensure
@@ -651,11 +649,14 @@ def _encode_datetime_with_cftime(dates, units: str, calendar: str) -> np.ndarray
651649
except TypeError:
652650
encoded_nums = cftime.date2num(filtered_dates, units, calendar)
653651

652+
if filtered_dates.size == none_position.size:
653+
return encoded_nums.reshape(original_shape)
654+
654655
# Create a full matrix of NaN
655656
# And fill the num dates in the not NaN or None position
656657
result = np.full(dates.shape, np.nan)
657658
result[np.nonzero(~none_position)] = encoded_nums
658-
return result
659+
return result.reshape(original_shape)
659660

660661

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

0 commit comments

Comments
 (0)