Merge pull request #1079 from jthielen/vendor-xarray

Vendor internal imports from xarray.core in xarray.py

metpy/_vendor/__init__.py

@@ -0,0 +1,4 @@
+# Copyright (c) 2019 MetPy Developers.
+# Distributed under the terms of the BSD 3-Clause License.
+# SPDX-License-Identifier: BSD-3-Clause
+"""Vendor external functionality used in MetPy."""

metpy/_vendor/xarray.py

@@ -0,0 +1,72 @@
+# Copyright (c) 2019 MetPy Developers.
+# Distributed under the terms of the BSD 3-Clause License.
+# SPDX-License-Identifier: BSD-3-Clause
+"""Vendor core functionality used from xarray.
+
+This code has been reproduced with modification under the terms of the Apache License, Version
+2.0 (notice included below).
+
+    Copyright 2014-2019, xarray Developers
+
+    Licensed under the Apache License, Version 2.0 (the "License");
+    you may not use this file except in compliance with the License.
+    You may obtain a copy of the License at
+
+        http://www.apache.org/licenses/LICENSE-2.0
+
+    Unless required by applicable law or agreed to in writing, software
+    distributed under the License is distributed on an "AS IS" BASIS,
+    WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+    See the License for the specific language governing permissions and
+    limitations under the License.
+"""
+
+
+def expanded_indexer(key, ndim):
+    """Expand an indexer to a tuple with length ndim.
+
+    Given a key for indexing an ndarray, return an equivalent key which is a
+    tuple with length equal to the number of dimensions.
+    The expansion is done by replacing all `Ellipsis` items with the right
+    number of full slices and then padding the key with full slices so that it
+    reaches the appropriate dimensionality.
+    """
+    if not isinstance(key, tuple):
+        # numpy treats non-tuple keys equivalent to tuples of length 1
+        key = (key,)
+    new_key = []
+    # handling Ellipsis right is a little tricky, see:
+    # http://docs.scipy.org/doc/numpy/reference/arrays.indexing.html#advanced-indexing
+    found_ellipsis = False
+    for k in key:
+        if k is Ellipsis:
+            if not found_ellipsis:
+                new_key.extend((ndim + 1 - len(key)) * [slice(None)])
+                found_ellipsis = True
+            else:
+                new_key.append(slice(None))
+        else:
+            new_key.append(k)
+    if len(new_key) > ndim:
+        raise IndexError('too many indices')
+    new_key.extend((ndim - len(new_key)) * [slice(None)])
+    return tuple(new_key)
+
+
+def is_dict_like(value):
+    """Check if value is dict-like."""
+    return hasattr(value, 'keys') and hasattr(value, '__getitem__')
+
+
+def either_dict_or_kwargs(pos_kwargs, kw_kwargs, func_name):
+    """Ensure dict-like argument from either positional or keyword arguments."""
+    if pos_kwargs is not None:
+        if not is_dict_like(pos_kwargs):
+            raise ValueError('the first argument to .{} must be a '
+                             'dictionary'.format(func_name))
+        if kw_kwargs:
+            raise ValueError('cannot specify both keyword and positional arguments to '
+                             '.{}'.format(func_name))
+        return pos_kwargs
+    else:
+        return kw_kwargs

metpy/tests/test_xarray.py

@@ -507,6 +507,24 @@ def test_data_array_loc_set_with_units(test_var):
     assert not np.isnan(temperature.loc[:, 700.]).any()
 
 
+def test_data_array_loc_with_ellipsis(test_var):
+    """Test the .loc indexer using multiple Ellipses to verify expansion behavior."""
+    truth = test_var[:, :, -1, :]
+    assert truth.identical(test_var.metpy.loc[..., 711.3653535503963 * units.km, ...])
+
+
+def test_data_array_loc_non_tuple(test_var):
+    """Test the .loc indexer with a non-tuple indexer."""
+    truth = test_var[-1]
+    assert truth.identical(test_var.metpy.loc['1987-04-04T18:00'])
+
+
+def test_data_array_loc_too_many_indices(test_var):
+    """Test the .loc indexer when too many indices are given."""
+    with pytest.raises(IndexError):
+        test_var.metpy.loc[:, 8.5e4 * units.Pa, :, :, :]
+
+
 def test_data_array_sel_dict_with_units(test_var):
     """Test .sel on the metpy accessor with dictionary."""
     truth = test_var.squeeze().loc[500.]
@@ -538,6 +556,21 @@ def test_dataset_sel_kwargs_with_units(test_ds):
                                              method='nearest'))
 
 
+def test_dataset_sel_non_dict_pos_arg(test_ds):
+    """Test that .sel errors when first positional argument is not a dict."""
+    with pytest.raises(ValueError) as exc:
+        test_ds.metpy.sel('1987-04-04T18:00:00')
+    assert 'must be a dictionary' in str(exc.value)
+
+
+def test_dataset_sel_mixed_dict_and_kwarg(test_ds):
+    """Test that .sel errors when dict positional argument and kwargs are mixed."""
+    with pytest.raises(ValueError) as exc:
+        test_ds.metpy.sel({'isobaric': slice(8.5e4 * units.Pa, 5e4 * units.Pa)},
+                          time='1987-04-04T18:00:00')
+    assert 'cannot specify both keyword and positional arguments' in str(exc.value)
+
+
 def test_dataset_loc_without_dict(test_ds):
     """Test that .metpy.loc for Datasets raises error when used with a non-dict."""
     with pytest.raises(TypeError):

metpy/xarray.py

@@ -10,9 +10,8 @@
 import warnings
 
 import xarray as xr
-from xarray.core.indexing import expanded_indexer
-from xarray.core.utils import either_dict_or_kwargs, is_dict_like
 
+from ._vendor.xarray import either_dict_or_kwargs, expanded_indexer, is_dict_like
 from .units import DimensionalityError, units
 
 __all__ = []