Skip to content

Commit b942724

Browse files
committed
First pass at making NetCDFTimeIndex compatible with pydata#1356
1 parent 2a7b439 commit b942724

File tree

4 files changed

+49
-10
lines changed

4 files changed

+49
-10
lines changed

xarray/conventions/netcdftimeindex.py

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -101,7 +101,7 @@ def _parsed_string_to_bounds(date_type, resolution, parsed):
101101

102102
def get_date_field(datetimes, field):
103103
"""Adapted from pandas.tslib.get_date_field"""
104-
return [getattr(date, field) for date in datetimes]
104+
return np.array([getattr(date, field) for date in datetimes])
105105

106106

107107
def _field_accessor(name, docstring=None):

xarray/core/accessors.py

Lines changed: 27 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -2,7 +2,7 @@
22
from __future__ import division
33
from __future__ import print_function
44

5-
from .common import is_datetime_like
5+
from .common import is_np_datetime_like, _contains_datetime_like_objects
66
from .pycompat import dask_array_type
77

88
from functools import partial
@@ -20,6 +20,20 @@ def _season_from_months(months):
2020
return seasons[(months // 3) % 4]
2121

2222

23+
def _access_through_netcdftimeindex(values, name):
24+
"""Coerce an array of datetime-like values to a NetCDFTimeIndex
25+
and access requested datetime component
26+
"""
27+
from ..conventions.netcdftimeindex import NetCDFTimeIndex
28+
values_as_netcdftimeindex = NetCDFTimeIndex(values)
29+
if name == 'season':
30+
months = values_as_netcdftimeindex.month
31+
field_values = _season_from_months(months)
32+
else:
33+
field_values = getattr(values_as_netcdftimeindex, name)
34+
return field_values.reshape(values.shape)
35+
36+
2337
def _access_through_series(values, name):
2438
"""Coerce an array of datetime-like values to a pandas Series and
2539
access requested datetime component
@@ -52,12 +66,17 @@ def _get_date_field(values, name, dtype):
5266
Array-like of datetime fields accessed for each element in values
5367
5468
"""
69+
if is_np_datetime_like(values.dtype):
70+
access_method = _access_through_series
71+
else:
72+
access_method = _access_through_netcdftimeindex
73+
5574
if isinstance(values, dask_array_type):
5675
from dask.array import map_blocks
57-
return map_blocks(_access_through_series,
76+
return map_blocks(access_method,
5877
values, name, dtype=dtype)
5978
else:
60-
return _access_through_series(values, name)
79+
return access_method(values, name)
6180

6281

6382
class DatetimeAccessor(object):
@@ -83,9 +102,11 @@ class DatetimeAccessor(object):
83102
84103
"""
85104
def __init__(self, xarray_obj):
86-
if not is_datetime_like(xarray_obj.dtype):
105+
if not _contains_datetime_like_objects(xarray_obj):
87106
raise TypeError("'dt' accessor only available for "
88-
"DataArray with datetime64 or timedelta64 dtype")
107+
"DataArray with datetime64 timedelta64 dtype or "
108+
"for arrays containing netcdftime datetime "
109+
"objects.")
89110
self._obj = xarray_obj
90111

91112
def _tslib_field_accessor(name, docstring=None, dtype=None):
@@ -147,4 +168,4 @@ def f(self, dtype=dtype):
147168

148169
time = _tslib_field_accessor(
149170
"time", "Timestamps corresponding to datetimes", object
150-
)
171+
)

xarray/core/common.py

Lines changed: 18 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -763,8 +763,25 @@ def ones_like(other, dtype=None):
763763
return full_like(other, 1, dtype)
764764

765765

766-
def is_datetime_like(dtype):
766+
def is_np_datetime_like(dtype):
767767
"""Check if a dtype is a subclass of the numpy datetime types
768768
"""
769769
return (np.issubdtype(dtype, np.datetime64) or
770770
np.issubdtype(dtype, np.timedelta64))
771+
772+
773+
def _contains_netcdftime_datetimes(var):
774+
"""Check if a variable contains netcdftime datetime objects"""
775+
from netcdftime._netcdftime import datetime
776+
return isinstance(var.data.ravel()[0], datetime)
777+
778+
779+
def _contains_datetime_like_objects(var):
780+
"""Check if a variable contains datetime like objects (either
781+
np.datetime64, np.timedelta64, or netcdftime._netcdftime.datetime)"""
782+
if is_np_datetime_like(var.dtype):
783+
return True
784+
try:
785+
return _contains_netcdftime_datetimes(var)
786+
except ImportError:
787+
return False

xarray/core/dataset.py

Lines changed: 3 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -21,7 +21,8 @@
2121
from .alignment import align
2222
from ..conventions import coding
2323
from .coordinates import DatasetCoordinates, LevelCoordinatesSource, Indexes
24-
from .common import ImplementsDatasetReduce, BaseDataObject, is_datetime_like
24+
from .common import (ImplementsDatasetReduce, BaseDataObject,
25+
_contains_datetime_like_objects)
2526
from .merge import (dataset_update_method, dataset_merge_method,
2627
merge_data_and_coords)
2728
from .utils import (Frozen, SortedKeysDict, maybe_wrap_array, hashable,
@@ -76,7 +77,7 @@ def _get_virtual_variable(variables, key, level_vars=None, dim_sizes=None):
7677
virtual_var = ref_var
7778
var_name = key
7879
else:
79-
if is_datetime_like(ref_var.dtype):
80+
if _contains_datetime_like_objects(ref_var):
8081
ref_var = xr.DataArray(ref_var)
8182
data = getattr(ref_var.dt, var_name).data
8283
else:

0 commit comments

Comments
 (0)