6
6
7
7
import numpy as np
8
8
9
- from pandas ._libs import lib , tslibs
9
+ from pandas ._libs import iNaT , lib , tslibs
10
10
import pandas .compat as compat
11
11
12
12
from pandas .core .dtypes .cast import _int64_max , maybe_upcast_putmask
13
13
from pandas .core .dtypes .common import (
14
14
_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 ,
17
17
is_object_dtype , is_scalar , is_timedelta64_dtype )
18
18
from pandas .core .dtypes .missing import isna , na_value_for_dtype , notna
19
19
@@ -203,15 +203,28 @@ def _get_values(values, skipna, fill_value=None, fill_value_typ=None,
203
203
if necessary copy and mask using the specified fill_value
204
204
copy = True will force the copy
205
205
"""
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
207
215
208
216
if mask is None :
209
217
if isfinite :
210
218
mask = _isfinite (values )
211
219
else :
212
220
mask = isna (values )
213
221
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
+
215
228
dtype_ok = _na_ok_dtype (dtype )
216
229
217
230
# 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,
232
245
elif copy :
233
246
values = values .copy ()
234
247
235
- values = _view_if_needed (values )
236
-
237
248
# return a platform independent precision dtype
238
249
dtype_max = dtype
239
250
if is_integer_dtype (dtype ) or is_bool_dtype (dtype ):
@@ -259,21 +270,19 @@ def _na_ok_dtype(dtype):
259
270
(np .integer , np .timedelta64 , np .datetime64 ))
260
271
261
272
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
-
268
273
def _wrap_results (result , dtype , fill_value = None ):
269
274
""" wrap our results if needed """
270
275
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
272
280
if not isinstance (result , np .ndarray ):
281
+ tz = getattr (dtype , 'tz' , None )
273
282
assert not isna (fill_value ), "Expected non-null fill_value"
274
283
if result == fill_value :
275
284
result = np .nan
276
- result = tslibs .Timestamp (result )
285
+ result = tslibs .Timestamp (result , tz = tz )
277
286
else :
278
287
result = result .view (dtype )
279
288
elif is_timedelta64_dtype (dtype ):
@@ -426,7 +435,6 @@ def nansum(values, axis=None, skipna=True, min_count=0, mask=None):
426
435
return _wrap_results (the_sum , dtype )
427
436
428
437
429
- @disallow ('M8' )
430
438
@bottleneck_switch ()
431
439
def nanmean (values , axis = None , skipna = True , mask = None ):
432
440
"""
@@ -457,7 +465,8 @@ def nanmean(values, axis=None, skipna=True, mask=None):
457
465
values , skipna , 0 , mask = mask )
458
466
dtype_sum = dtype_max
459
467
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 )):
461
470
dtype_sum = np .float64
462
471
elif is_float_dtype (dtype ):
463
472
dtype_sum = dtype
@@ -466,7 +475,9 @@ def nanmean(values, axis=None, skipna=True, mask=None):
466
475
the_sum = _ensure_numeric (values .sum (axis , dtype = dtype_sum ))
467
476
468
477
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
470
481
ct_mask = count == 0
471
482
if ct_mask .any ():
472
483
the_mean [ct_mask ] = np .nan
0 commit comments