Update remaining accessor functionality to no longer be in-place

src/metpy/xarray.py

@@ -229,14 +229,17 @@ def cartopy_globe(self):
 
     def _fixup_coordinate_map(self, coord_map):
         """Ensure sure we have coordinate variables in map, not coordinate names."""
+        new_coord_map = {}
         for axis in coord_map:
             if coord_map[axis] is not None and not isinstance(coord_map[axis], xr.DataArray):
-                coord_map[axis] = self._data_array[coord_map[axis]]
+                new_coord_map[axis] = self._data_array[coord_map[axis]]
+            else:
+                new_coord_map[axis] = coord_map[axis]
 
-        return coord_map
+        return new_coord_map
 
     def assign_coordinates(self, coordinates):
-        """Assign the given coordinates to the given MetPy axis types.
+        """Return new DataArray with given coordinates assigned to the given MetPy axis types.
 
         Parameters
         ----------
@@ -247,18 +250,32 @@ def assign_coordinates(self, coordinates):
             which will trigger reparsing of all coordinates on next access.
 
         """
+        coord_updates = {}
         if coordinates:
             # Assign the _metpy_axis attributes according to supplied mapping
             coordinates = self._fixup_coordinate_map(coordinates)
             for axis in coordinates:
                 if coordinates[axis] is not None:
-                    _assign_axis(coordinates[axis].attrs, axis)
+                    coord_updates[coordinates[axis].name] = (
+                        coordinates[axis].assign_attrs(
+                            _assign_axis(coordinates[axis].attrs.copy(), axis)
+                        )
+                    )
         else:
             # Clear _metpy_axis attribute on all coordinates
-            for coord_var in self._data_array.coords.values():
-                coord_var.attrs.pop('_metpy_axis', None)
+            for coord_name, coord_var in self._data_array.coords.items():
+                coord_updates[coord_name] = coord_var.copy(deep=False)
+
+                # Some coordinates remained linked in old form under other coordinates. We
+                # need to remove from these.
+                sub_coords = coord_updates[coord_name].coords
+                for sub_coord in sub_coords:
+                    coord_updates[coord_name].coords[sub_coord].attrs.pop('_metpy_axis', None)
+
+                # Now we can remove the _metpy_axis attr from the coordinate itself
+                coord_updates[coord_name].attrs.pop('_metpy_axis', None)
 
-        return self._data_array  # allow method chaining
+        return self._data_array.assign_coords(coord_updates)
 
     def _generate_coordinate_map(self):
         """Generate a coordinate map via CF conventions and other methods."""
@@ -317,6 +334,11 @@ def _metpy_axis_search(self, metpy_axis):
                 return coord_var
 
         # Opportunistically parse all coordinates, and assign if not already assigned
+        # Note: since this is generally called by way of the coordinate properties, to cache
+        # the coordinate parsing results in coord_map on the coordinates means modifying the
+        # DataArray in-place (an exception to the usual behavior of MetPy's accessor). This is
+        # considered safe because it only effects the "_metpy_axis" attribute on the
+        # coordinates, and nothing else.
         coord_map = self._generate_coordinate_map()
         for axis, coord_var in coord_map.items():
             if (coord_var is not None
@@ -651,7 +673,7 @@ def parse_cf(self, varname=None, coordinates=None):
 
         # Assign coordinates if the coordinates argument is given
         if coordinates is not None:
-            var.metpy.assign_coordinates(coordinates)
+            var = var.metpy.assign_coordinates(coordinates)
 
         # Attempt to build the crs coordinate
         crs = None
@@ -840,7 +862,7 @@ def assign_y_x(self, force=False, tolerance=None):
             return self._dataset.assign_coords(**{y.name: y, x.name: x})
 
     def update_attribute(self, attribute, mapping):
-        """Update attribute of all Dataset variables.
+        """Return new Dataset with specified attribute updated on all Dataset variables.
 
         Parameters
         ----------
@@ -855,24 +877,33 @@ def update_attribute(self, attribute, mapping):
         Returns
         -------
         `xarray.Dataset`
-            Dataset with attribute updated (modified in place, and returned to allow method
-            chaining)
+            New Dataset with attribute updated
 
         """
         # Make mapping uniform
-        if callable(mapping):
-            mapping_func = mapping
-        else:
-            def mapping_func(varname, **kwargs):
-                return mapping.get(varname, None)
+        if not callable(mapping):
+            old_mapping = mapping
 
-        # Apply across all variables
-        for varname in list(self._dataset.data_vars) + list(self._dataset.coords):
-            value = mapping_func(varname, **self._dataset[varname].attrs)
-            if value is not None:
-                self._dataset[varname].attrs[attribute] = value
+            def mapping(varname, **kwargs):
+                return old_mapping.get(varname, None)
 
-        return self._dataset
+        # Define mapping function for Dataset.map
+        def mapping_func(da):
+            new_value = mapping(da.name, **da.attrs)
+            if new_value is None:
+                return da
+            else:
+                return da.assign_attrs(**{attribute: new_value})
+
+        # Apply across all variables and coordinates
+        return (
+            self._dataset
+            .map(mapping_func, keep_attrs=True)
+            .assign_coords({
+                coord_name: mapping_func(coord_var)
+                for coord_name, coord_var in self._dataset.coords.items()
+            })
+        )
 
     def quantify(self):
         """Return new dataset with all numeric variables quantified and cached data loaded."""

tests/test_xarray.py

@@ -302,7 +302,7 @@ def test_assign_coordinates_not_overwrite(test_ds_generic):
     """Test that assign_coordinates does not overwrite past axis attributes."""
     data = test_ds_generic.copy()
     data['c'].attrs['axis'] = 'X'
-    data['test'].metpy.assign_coordinates({'y': data['c']})
+    data['test'] = data['test'].metpy.assign_coordinates({'y': data['c']})
     assert data['c'].identical(data['test'].metpy.y)
     assert data['c'].attrs['axis'] == 'X'
 
@@ -622,9 +622,12 @@ def test_data_array_sel_dict_with_units(test_var):
 def test_data_array_sel_kwargs_with_units(test_var):
     """Test .sel on the metpy accessor with kwargs and axis type."""
     truth = test_var.loc[:, 500.][..., 122]
-    selection = test_var.metpy.sel(vertical=5e4 * units.Pa, x=-16.569 * units.km,
-                                   tolerance=1., method='nearest')
-    selection.metpy.assign_coordinates(None)  # truth was not parsed for coordinates
+    selection = (
+        test_var.metpy
+        .sel(vertical=5e4 * units.Pa, x=-16.569 * units.km, tolerance=1., method='nearest')
+        .metpy
+        .assign_coordinates(None)
+    )
     assert truth.identical(selection)
 
 
@@ -983,24 +986,37 @@ def test_update_attribute_dictionary(test_ds_generic):
         'test': 'Filler data',
         'c': 'The third coordinate'
     }
-    test_ds_generic.metpy.update_attribute('description', descriptions)
-    assert 'description' not in test_ds_generic['a'].attrs
-    assert 'description' not in test_ds_generic['b'].attrs
-    assert test_ds_generic['c'].attrs['description'] == 'The third coordinate'
-    assert 'description' not in test_ds_generic['d'].attrs
-    assert 'description' not in test_ds_generic['e'].attrs
-    assert test_ds_generic['test'].attrs['description'] == 'Filler data'
+    result = test_ds_generic.metpy.update_attribute('description', descriptions)
+
+    # Test attribute updates
+    assert 'description' not in result['a'].attrs
+    assert 'description' not in result['b'].attrs
+    assert result['c'].attrs['description'] == 'The third coordinate'
+    assert 'description' not in result['d'].attrs
+    assert 'description' not in result['e'].attrs
+    assert result['test'].attrs['description'] == 'Filler data'
+
+    # Test for no side effects
+    assert 'description' not in test_ds_generic['c'].attrs
+    assert 'description' not in test_ds_generic['test'].attrs
 
 
 def test_update_attribute_callable(test_ds_generic):
     """Test update_attribute using callable."""
     def even_ascii(varname, **kwargs):
         if ord(varname[0]) % 2 == 0:
             return 'yes'
-    test_ds_generic.metpy.update_attribute('even', even_ascii)
-    assert 'even' not in test_ds_generic['a'].attrs
-    assert test_ds_generic['b'].attrs['even'] == 'yes'
-    assert 'even' not in test_ds_generic['c'].attrs
-    assert test_ds_generic['d'].attrs['even'] == 'yes'
-    assert 'even' not in test_ds_generic['e'].attrs
-    assert test_ds_generic['test'].attrs['even'] == 'yes'
+    result = test_ds_generic.metpy.update_attribute('even', even_ascii)
+
+    # Test attribute updates
+    assert 'even' not in result['a'].attrs
+    assert result['b'].attrs['even'] == 'yes'
+    assert 'even' not in result['c'].attrs
+    assert result['d'].attrs['even'] == 'yes'
+    assert 'even' not in result['e'].attrs
+    assert result['test'].attrs['even'] == 'yes'
+
+    # Test for no side effects
+    assert 'even' not in test_ds_generic['b'].attrs
+    assert 'even' not in test_ds_generic['d'].attrs
+    assert 'even' not in test_ds_generic['test'].attrs