Add a filter_by_attrs method to Dataset (#844)

* get_variables_by_attributes

* tests

* Changelog entry

* Review actions

* Improved whats_new

* Changing to data variables only

* Return an empty Dataset

* Review actions

* More review actions

* Fix docstring indentation

* Create an example dataset

* api docs

* Rename to filter_by_attrs

Author: ocefpaf

doc/api.rst

@@ -395,6 +395,7 @@ Dataset methods
    Dataset.close
    Dataset.load
    Dataset.chunk
+   Dataset.filter_by_attrs
 
 DataArray methods
 -----------------

doc/whats-new.rst

@@ -80,6 +80,11 @@ Enhancements
   allowing more control on the colorbar (:issue:`872`).
   By `Fabien Maussion <https://github.com/fmaussion>`_.
 
+- New Dataset method :py:meth:`filter_by_attrs`, akin to
+  ``netCDF4.Dataset.get_variables_by_attributes``, to easily filter
+  data variables using its attributes.
+  `Filipe Fernandes <https://github.com/ocefpaf>`_.
+
 Bug fixes
 ~~~~~~~~~
 

xarray/core/dataset.py

@@ -2176,5 +2176,88 @@ def real(self):
     def imag(self):
         return self._unary_op(lambda x: x.imag, keep_attrs=True)(self)
 
+    def filter_by_attrs(self, **kwargs):
+        """Returns a ``Dataset`` with variables that match specific conditions.
+
+        Can pass in ``key=value ``or ``key=callable``.  Variables are returned
+        that contain all of the matches or callable returns True.  If using a
+        callable note that it should accept a single parameter only,
+        the attribute value.
+
+        Parameters
+        ----------
+        **kwargs : key=value
+            key : str
+                Attribute name.
+            value : callable or obj
+                If value is a callable, it should return a boolean in the form
+                of bool = func(attr) where attr is da.attrs[key].
+                Otherwise, value will be compared to the each
+                DataArray's attrs[key].
+
+        Returns
+        -------
+        new : Dataset
+            New dataset with variables filtered by attribute.
+
+        Examples
+        --------
+        # "Create an example dataset:
+        >>> import numpy as np
+        >>> import pandas as pd
+        >>> import xarray as xr
+        >>> temp = 15 + 8 * np.random.randn(2, 2, 3)
+        >>> precip = 10 * np.random.rand(2, 2, 3)
+        >>> lon = [[-99.83, -99.32], [-99.79, -99.23]]
+        >>> lat = [[42.25, 42.21], [42.63, 42.59]]
+        >>> dims = ['x', 'y', 'time']
+        >>> temp_attr = dict(standard_name='air_potential_temperature')
+        >>> precip_attr = dict(standard_name='convective_precipitation_flux')
+        >>> ds = xr.Dataset({
+        ...         'temperature': (dims,  temp, temp_attr),
+        ...         'precipitation': (dims, precip, precip_attr)},
+        ...                 coords={
+        ...         'lon': (['x', 'y'], lon),
+        ...         'lat': (['x', 'y'], lat),
+        ...         'time': pd.date_range('2014-09-06', periods=3),
+        ...         'reference_time': pd.Timestamp('2014-09-05')})
+        >>> # Get variables matching a specific standard_name.
+        >>> ds.filter_by_attrs(standard_name='convective_precipitation_flux')
+        <xarray.Dataset>
+        Dimensions:         (time: 3, x: 2, y: 2)
+        Coordinates:
+          * x               (x) int64 0 1
+          * time            (time) datetime64[ns] 2014-09-06 2014-09-07 2014-09-08
+            lat             (x, y) float64 42.25 42.21 42.63 42.59
+          * y               (y) int64 0 1
+            reference_time  datetime64[ns] 2014-09-05
+            lon             (x, y) float64 -99.83 -99.32 -99.79 -99.23
+        Data variables:
+            precipitation   (x, y, time) float64 4.178 2.307 6.041 6.046 0.06648 ...
+        >>> # Get all variables that have a standard_name attribute.
+        >>> standard_name = lambda v: v is not None
+        >>> ds.filter_by_attrs(standard_name=standard_name)
+        <xarray.Dataset>
+        Dimensions:         (time: 3, x: 2, y: 2)
+        Coordinates:
+            lon             (x, y) float64 -99.83 -99.32 -99.79 -99.23
+            lat             (x, y) float64 42.25 42.21 42.63 42.59
+          * x               (x) int64 0 1
+          * y               (y) int64 0 1
+          * time            (time) datetime64[ns] 2014-09-06 2014-09-07 2014-09-08
+            reference_time  datetime64[ns] 2014-09-05
+        Data variables:
+            temperature     (x, y, time) float64 25.86 20.82 6.954 23.13 10.25 11.68 ...
+            precipitation   (x, y, time) float64 5.702 0.9422 2.075 1.178 3.284 ...
+
+        """
+        selection = []
+        for var_name, variable in self.data_vars.items():
+            for attr_name, pattern in kwargs.items():
+                attr_value = variable.attrs.get(attr_name)
+                if ((callable(pattern) and pattern(attr_value))
+                        or attr_value == pattern):
+                    selection.append(var_name)
+        return self[selection]
 
 ops.inject_all_ops_and_reduce_methods(Dataset, array_only=False)

xarray/test/test_dataset.py

@@ -2494,3 +2494,39 @@ def test_setattr_raises(self):
             ds.foo = 2
         with self.assertRaisesRegexp(AttributeError, 'cannot set attr'):
             ds.other = 2
+
+    def test_filter_by_attrs(self):
+        precip = dict(standard_name='convective_precipitation_flux')
+        temp0 = dict(standard_name='air_potential_temperature', height='0 m')
+        temp10 = dict(standard_name='air_potential_temperature', height='10 m')
+        ds = Dataset({'temperature_0': (['t'], [0], temp0),
+                      'temperature_10': (['t'], [0], temp10),
+                      'precipitation': (['t'], [0], precip)},
+                    coords={'time': (['t'], [0], dict(axis='T'))})
+
+        # Test return empty Dataset.
+        ds.filter_by_attrs(standard_name='invalid_standard_name')
+        new_ds = ds.filter_by_attrs(standard_name='invalid_standard_name')
+        self.assertFalse(bool(new_ds.data_vars))
+
+        # Test return one DataArray.
+        new_ds = ds.filter_by_attrs(standard_name='convective_precipitation_flux')
+        self.assertEqual(new_ds['precipitation'].standard_name, 'convective_precipitation_flux')
+        self.assertDatasetEqual(new_ds['precipitation'], ds['precipitation'])
+
+        # Test return more than one DataArray.
+        new_ds = ds.filter_by_attrs(standard_name='air_potential_temperature')
+        self.assertEqual(len(new_ds.data_vars), 2)
+        for var in new_ds.data_vars:
+            self.assertEqual(new_ds[var].standard_name, 'air_potential_temperature')
+
+        # Test callable.
+        new_ds = ds.filter_by_attrs(height=lambda v: v is not None)
+        self.assertEqual(len(new_ds.data_vars), 2)
+        for var in new_ds.data_vars:
+            self.assertEqual(new_ds[var].standard_name, 'air_potential_temperature')
+
+        new_ds = ds.filter_by_attrs(height='10 m')
+        self.assertEqual(len(new_ds.data_vars), 1)
+        for var in new_ds.data_vars:
+            self.assertEqual(new_ds[var].height, '10 m')