66
77import  numpy  as  np 
88
9- from  pandas ._libs  import  lib , tslibs 
9+ from  pandas ._libs  import  iNaT ,  lib , tslibs 
1010import  pandas .compat  as  compat 
1111
1212from  pandas .core .dtypes .cast  import  _int64_max , maybe_upcast_putmask 
1313from  pandas .core .dtypes .common  import  (
1414    _get_dtype , is_any_int_dtype , is_bool_dtype , is_complex , is_complex_dtype ,
15-     is_datetime64_dtype , is_datetime_or_timedelta_dtype ,  is_float ,
16-     is_float_dtype , is_integer , is_integer_dtype , is_numeric_dtype ,
15+     is_datetime64_dtype , is_datetime64tz_dtype ,  is_datetime_or_timedelta_dtype ,
16+     is_float ,  is_float_dtype , is_integer , is_integer_dtype , is_numeric_dtype ,
1717    is_object_dtype , is_scalar , is_timedelta64_dtype )
1818from  pandas .core .dtypes .missing  import  isna , na_value_for_dtype , notna 
1919
@@ -203,15 +203,28 @@ def _get_values(values, skipna, fill_value=None, fill_value_typ=None,
203203    if necessary copy and mask using the specified fill_value 
204204    copy = True will force the copy 
205205    """ 
206-     values  =  com .values_from_object (values )
206+ 
207+     if  is_datetime64tz_dtype (values ):
208+         # com.values_from_object returns M8[ns] dtype instead of tz-aware, 
209+         #  so this case must be handled separately from the rest 
210+         dtype  =  values .dtype 
211+         values  =  getattr (values , "_values" , values )
212+     else :
213+         values  =  com .values_from_object (values )
214+         dtype  =  values .dtype 
207215
208216    if  mask  is  None :
209217        if  isfinite :
210218            mask  =  _isfinite (values )
211219        else :
212220            mask  =  isna (values )
213221
214-     dtype  =  values .dtype 
222+     if  is_datetime_or_timedelta_dtype (values ) or  is_datetime64tz_dtype (values ):
223+         # changing timedelta64/datetime64 to int64 needs to happen after 
224+         #  finding `mask` above 
225+         values  =  getattr (values , "asi8" , values )
226+         values  =  values .view (np .int64 )
227+ 
215228    dtype_ok  =  _na_ok_dtype (dtype )
216229
217230    # get our fill value (in case we need to provide an alternative 
@@ -232,8 +245,6 @@ def _get_values(values, skipna, fill_value=None, fill_value_typ=None,
232245    elif  copy :
233246        values  =  values .copy ()
234247
235-     values  =  _view_if_needed (values )
236- 
237248    # return a platform independent precision dtype 
238249    dtype_max  =  dtype 
239250    if  is_integer_dtype (dtype ) or  is_bool_dtype (dtype ):
@@ -259,21 +270,19 @@ def _na_ok_dtype(dtype):
259270                          (np .integer , np .timedelta64 , np .datetime64 ))
260271
261272
262- def  _view_if_needed (values ):
263-     if  is_datetime_or_timedelta_dtype (values ):
264-         return  values .view (np .int64 )
265-     return  values 
266- 
267- 
268273def  _wrap_results (result , dtype , fill_value = None ):
269274    """ wrap our results if needed """ 
270275
271-     if  is_datetime64_dtype (dtype ):
276+     if  is_datetime64_dtype (dtype ) or  is_datetime64tz_dtype (dtype ):
277+         if  fill_value  is  None :
278+             # GH#24293 
279+             fill_value  =  iNaT 
272280        if  not  isinstance (result , np .ndarray ):
281+             tz  =  getattr (dtype , 'tz' , None )
273282            assert  not  isna (fill_value ), "Expected non-null fill_value" 
274283            if  result  ==  fill_value :
275284                result  =  np .nan 
276-             result  =  tslibs .Timestamp (result )
285+             result  =  tslibs .Timestamp (result ,  tz = tz )
277286        else :
278287            result  =  result .view (dtype )
279288    elif  is_timedelta64_dtype (dtype ):
@@ -426,7 +435,6 @@ def nansum(values, axis=None, skipna=True, min_count=0, mask=None):
426435    return  _wrap_results (the_sum , dtype )
427436
428437
429- @disallow ('M8' ) 
430438@bottleneck_switch () 
431439def  nanmean (values , axis = None , skipna = True , mask = None ):
432440    """ 
@@ -457,7 +465,8 @@ def nanmean(values, axis=None, skipna=True, mask=None):
457465        values , skipna , 0 , mask = mask )
458466    dtype_sum  =  dtype_max 
459467    dtype_count  =  np .float64 
460-     if  is_integer_dtype (dtype ) or  is_timedelta64_dtype (dtype ):
468+     if  (is_integer_dtype (dtype ) or  is_timedelta64_dtype (dtype ) or 
469+             is_datetime64_dtype (dtype ) or  is_datetime64tz_dtype (dtype )):
461470        dtype_sum  =  np .float64 
462471    elif  is_float_dtype (dtype ):
463472        dtype_sum  =  dtype 
@@ -466,7 +475,9 @@ def nanmean(values, axis=None, skipna=True, mask=None):
466475    the_sum  =  _ensure_numeric (values .sum (axis , dtype = dtype_sum ))
467476
468477    if  axis  is  not   None  and  getattr (the_sum , 'ndim' , False ):
469-         the_mean  =  the_sum  /  count 
478+         with  np .errstate (all = "ignore" ):
479+             # suppress division by zero warnings 
480+             the_mean  =  the_sum  /  count 
470481        ct_mask  =  count  ==  0 
471482        if  ct_mask .any ():
472483            the_mean [ct_mask ] =  np .nan 
0 commit comments