@@ -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
649661def cast_to_int_if_safe (num ) -> np .ndarray :
0 commit comments