Repository: pydata/xarray

Issue:
Bug in legend of dataset.plot.scatter
<!-- A short summary of the issue, if appropriate -->

When using Dataset.scatter with hue being a variable of dtype string, the legend turns out to be wrong.
#### MCVE Code Sample
```python
import xarray as xr
import numpy as np
dd = xr.Dataset({'y': (['x'], np.arange(8)),
                 'label': (['x'], list('AABBCCDD'))},
                coords={'x': np.linspace(0,1,8)})
dd.plot.scatter(x='x', y='y', hue='label')
```
Output is (note the legend):
![Figure_1](https://user-images.githubusercontent.com/6164157/83973237-a9572e80-a8ed-11ea-8aea-642dce259b2b.png)

Playing around it seems that it always chooses the first 4 values as the legend labels (note that the order of colors of the points is correct):
```python
import xarray as xr
import numpy as np
dd = xr.Dataset({'y': (['x'], np.arange(8)),
                 'label': (['x'], list('ABBACDDC'))},
                coords={'x': np.linspace(0,1,8)})
dd.plot.scatter(x='x', y='y', hue='label')
```
![Figure_2](https://user-images.githubusercontent.com/6164157/83973242-b96f0e00-a8ed-11ea-899d-72f0f82baf8a.png)

And if there are only 3 labels in total it chooses the first 3:
```python
import xarray as xr
import numpy as np
dd = xr.Dataset({'y': (['x'], np.arange(6)),
                 'label': (['x'], list('ABBACC'))},
                coords={'x': np.linspace(0,1,6)})
dd.plot.scatter(x='x', y='y', hue='label')
```
![Figure_1](https://user-images.githubusercontent.com/6164157/83973353-8e38ee80-a8ee-11ea-9c75-807fc2558dfe.png)

#### Expected Output
Legend in first two plots should read 'ABCD' and last plot 'ABC'

#### Versions

<details><summary>Output of <tt>xr.show_versions()</tt></summary>


INSTALLED VERSIONS
------------------
commit: None
libhdf5: 1.10.4
libnetcdf: None

xarray: 0.15.1
pandas: 0.25.1
numpy: 1.17.2
scipy: 1.3.1
netCDF4: None
pydap: None
h5netcdf: None
h5py: 2.9.0
Nio: None
zarr: None
cftime: None
nc_time_axis: None
PseudoNetCDF: None
rasterio: None
cfgrib: None
iris: None
bottleneck: 1.2.1
dask: 2.5.2
distributed: 2.5.2
matplotlib: 3.1.1
cartopy: None
seaborn: 0.9.0
numbagg: None
setuptools: 41.4.0
pip: 19.2.3
conda: 4.8.1
pytest: 5.2.1
IPython: 7.8.0
sphinx: 2.2.0


</details>


Discussion:

Patch:
diff --git a/.pre-commit-config.yaml b/.pre-commit-config.yaml
index 7265bbac..16165288 100644
--- a/.pre-commit-config.yaml
+++ b/.pre-commit-config.yaml
@@ -6,7 +6,7 @@ repos:
     hooks:
       - id: isort
   # https://github.com/python/black#version-control-integration
-  - repo: https://github.com/python/black
+  - repo: https://github.com/psf/black
     rev: 20.8b1
     hooks:
       - id: black
diff --git a/doc/whats-new.rst b/doc/whats-new.rst
index 316254a9..74619529 100644
--- a/doc/whats-new.rst
+++ b/doc/whats-new.rst
@@ -77,12 +77,15 @@ Bug fixes
   and :py:meth:`DataArray.str.wrap` (:issue:`4334`). By `Mathias Hauser <https://github.com/mathause>`_.
 - Fixed overflow issue causing incorrect results in computing means of :py:class:`cftime.datetime`
   arrays (:issue:`4341`). By `Spencer Clark <https://github.com/spencerkclark>`_.
+- Fixed :py:meth:`Dataset.coarsen`, :py:meth:`DataArray.coarsen` dropping attributes on original object (:issue:`4120`, :pull:`4360`). by `Julia Kent <https://github.com/jukent>`_.
 - fix the signature of the plot methods. (:pull:`4359`) By `Justus Magin <https://github.com/keewis>`_.
 - Fix :py:func:`xarray.apply_ufunc` with ``vectorize=True`` and ``exclude_dims`` (:issue:`3890`).
   By `Mathias Hauser <https://github.com/mathause>`_.
 - Fix `KeyError` when doing linear interpolation to an nd `DataArray`
   that contains NaNs (:pull:`4233`).
   By `Jens Svensmark <https://github.com/jenssss>`_
+- Fix incorrect legend labels for :py:meth:`Dataset.plot.scatter` (:issue:`4126`).
+  By `Peter Hausamann <https://github.com/phausamann>`_.
 
 Documentation
 ~~~~~~~~~~~~~
@@ -97,7 +100,7 @@ Documentation
 Internal Changes
 ~~~~~~~~~~~~~~~~
 - Use :py:func:`dask.array.apply_gufunc` instead of :py:func:`dask.array.blockwise` in
-  :py:func:`xarray.apply_ufunc` when using ``dask='parallelized'``. (:pull:`4060`, :pull:`4391`)
+  :py:func:`xarray.apply_ufunc` when using ``dask='parallelized'``. (:pull:`4060`, :pull:`4391`, :pull:`4392`)
 - Fix ``pip install .`` when no ``.git`` directory exists; namely when the xarray source
   directory has been rsync'ed by PyCharm Professional for a remote deployment over SSH.
   By `Guido Imperiale <https://github.com/crusaderky>`_
diff --git a/xarray/core/computation.py b/xarray/core/computation.py
index 507f12fe..db4a9b46 100644
--- a/xarray/core/computation.py
+++ b/xarray/core/computation.py
@@ -625,6 +625,26 @@ def apply_variable_ufunc(
             if dask_gufunc_kwargs is None:
                 dask_gufunc_kwargs = {}
 
+            allow_rechunk = dask_gufunc_kwargs.get("allow_rechunk", None)
+            if allow_rechunk is None:
+                for n, (data, core_dims) in enumerate(
+                    zip(input_data, signature.input_core_dims)
+                ):
+                    if is_duck_dask_array(data):
+                        # core dimensions cannot span multiple chunks
+                        for axis, dim in enumerate(core_dims, start=-len(core_dims)):
+                            if len(data.chunks[axis]) != 1:
+                                raise ValueError(
+                                    f"dimension {dim} on {n}th function argument to "
+                                    "apply_ufunc with dask='parallelized' consists of "
+                                    "multiple chunks, but is also a core dimension. To "
+                                    "fix, either rechunk into a single dask array chunk along "
+                                    f"this dimension, i.e., ``.chunk({dim}: -1)``, or "
+                                    "pass ``allow_rechunk=True`` in ``dask_gufunc_kwargs`` "
+                                    "but beware that this may significantly increase memory usage."
+                                )
+                dask_gufunc_kwargs["allow_rechunk"] = True
+
             output_sizes = dask_gufunc_kwargs.pop("output_sizes", {})
             if output_sizes:
                 output_sizes_renamed = {}
diff --git a/xarray/core/rolling.py b/xarray/core/rolling.py
index 0c4614e0..3d7be955 100644
--- a/xarray/core/rolling.py
+++ b/xarray/core/rolling.py
@@ -690,7 +690,7 @@ class DataArrayCoarsen(Coarsen):
             from .dataarray import DataArray
 
             reduced = self.obj.variable.coarsen(
-                self.windows, func, self.boundary, self.side, **kwargs
+                self.windows, func, self.boundary, self.side, self.keep_attrs, **kwargs
             )
             coords = {}
             for c, v in self.obj.coords.items():
@@ -703,6 +703,7 @@ class DataArrayCoarsen(Coarsen):
                             self.coord_func[c],
                             self.boundary,
                             self.side,
+                            self.keep_attrs,
                             **kwargs,
                         )
                     else:
diff --git a/xarray/core/variable.py b/xarray/core/variable.py
index ce6df528..075d7904 100644
--- a/xarray/core/variable.py
+++ b/xarray/core/variable.py
@@ -1993,7 +1993,9 @@ class Variable(
             ),
         )
 
-    def coarsen(self, windows, func, boundary="exact", side="left", **kwargs):
+    def coarsen(
+        self, windows, func, boundary="exact", side="left", keep_attrs=None, **kwargs
+    ):
         """
         Apply reduction function.
         """
@@ -2001,13 +2003,22 @@ class Variable(
         if not windows:
             return self.copy()
 
+        if keep_attrs is None:
+            keep_attrs = _get_keep_attrs(default=False)
+
+        if keep_attrs:
+            _attrs = self.attrs
+        else:
+            _attrs = None
+
         reshaped, axes = self._coarsen_reshape(windows, boundary, side)
         if isinstance(func, str):
             name = func
             func = getattr(duck_array_ops, name, None)
             if func is None:
                 raise NameError(f"{name} is not a valid method.")
-        return self._replace(data=func(reshaped, axis=axes, **kwargs))
+
+        return self._replace(data=func(reshaped, axis=axes, **kwargs), attrs=_attrs)
 
     def _coarsen_reshape(self, windows, boundary, side):
         """
@@ -2072,9 +2083,6 @@ class Variable(
             else:
                 shape.append(variable.shape[i])
 
-        keep_attrs = _get_keep_attrs(default=False)
-        variable.attrs = variable._attrs if keep_attrs else {}
-
         return variable.data.reshape(shape), tuple(axes)
 
     @property
diff --git a/xarray/plot/dataset_plot.py b/xarray/plot/dataset_plot.py
index 51ceff17..c6158163 100644
--- a/xarray/plot/dataset_plot.py
+++ b/xarray/plot/dataset_plot.py
@@ -339,11 +339,7 @@ def _dsplot(plotfunc):
             ax.set_ylabel(meta_data.get("ylabel"))
 
         if meta_data["add_legend"]:
-            ax.legend(
-                handles=primitive,
-                labels=list(meta_data["hue"].values),
-                title=meta_data.get("hue_label", None),
-            )
+            ax.legend(handles=primitive, title=meta_data.get("hue_label", None))
         if meta_data["add_colorbar"]:
             cbar_kwargs = {} if cbar_kwargs is None else cbar_kwargs
             if "label" not in cbar_kwargs:


Test patch:
diff --git a/xarray/tests/test_computation.py b/xarray/tests/test_computation.py
index 7cb755b6..63bedfaf 100644
--- a/xarray/tests/test_computation.py
+++ b/xarray/tests/test_computation.py
@@ -695,8 +695,7 @@ def test_apply_dask_parallelized_two_args():
     check(data_array, 0 * data_array)
     check(data_array, 0 * data_array[0])
     check(data_array[:, 0], 0 * data_array[0])
-    with raises_regex(ValueError, "with different chunksize present"):
-        check(data_array, 0 * data_array.compute())
+    check(data_array, 0 * data_array.compute())
 
 
 @requires_dask
@@ -710,7 +709,7 @@ def test_apply_dask_parallelized_errors():
     with raises_regex(ValueError, "at least one input is an xarray object"):
         apply_ufunc(identity, array, dask="parallelized")
 
-    # formerly from _apply_blockwise, now from dask.array.apply_gufunc
+    # formerly from _apply_blockwise, now from apply_variable_ufunc
     with raises_regex(ValueError, "consists of multiple chunks"):
         apply_ufunc(
             identity,
diff --git a/xarray/tests/test_dataarray.py b/xarray/tests/test_dataarray.py
index 8fd56446..a22ed58c 100644
--- a/xarray/tests/test_dataarray.py
+++ b/xarray/tests/test_dataarray.py
@@ -6186,6 +6186,35 @@ def test_isin(da):
     assert_equal(result, expected)
 
 
+def test_coarsen_keep_attrs():
+    _attrs = {"units": "test", "long_name": "testing"}
+
+    da = xr.DataArray(
+        np.linspace(0, 364, num=364),
+        dims="time",
+        coords={"time": pd.date_range("15/12/1999", periods=364)},
+        attrs=_attrs,
+    )
+
+    da2 = da.copy(deep=True)
+
+    # Test dropped attrs
+    dat = da.coarsen(time=3, boundary="trim").mean()
+    assert dat.attrs == {}
+
+    # Test kept attrs using dataset keyword
+    dat = da.coarsen(time=3, boundary="trim", keep_attrs=True).mean()
+    assert dat.attrs == _attrs
+
+    # Test kept attrs using global option
+    with xr.set_options(keep_attrs=True):
+        dat = da.coarsen(time=3, boundary="trim").mean()
+    assert dat.attrs == _attrs
+
+    # Test kept attrs in original object
+    xr.testing.assert_identical(da, da2)
+
+
 @pytest.mark.filterwarnings("error:Mean of empty slice")
 @pytest.mark.parametrize("da", (1, 2), indirect=True)
 def test_rolling_iter(da):
diff --git a/xarray/tests/test_dataset.py b/xarray/tests/test_dataset.py
index 16b1cc33..40e2bdfc 100644
--- a/xarray/tests/test_dataset.py
+++ b/xarray/tests/test_dataset.py
@@ -5956,6 +5956,8 @@ def test_coarsen_keep_attrs():
         attrs=_attrs,
     )
 
+    ds2 = ds.copy(deep=True)
+
     # Test dropped attrs
     dat = ds.coarsen(coord=5).mean()
     assert dat.attrs == {}
@@ -5969,6 +5971,9 @@ def test_coarsen_keep_attrs():
         dat = ds.coarsen(coord=5).mean()
     assert dat.attrs == _attrs
 
+    # Test kept attrs in original object
+    xr.testing.assert_identical(ds, ds2)
+
 
 def test_rolling_keep_attrs():
     _attrs = {"units": "test", "long_name": "testing"}
diff --git a/xarray/tests/test_plot.py b/xarray/tests/test_plot.py
index 1dfbb97d..14677b40 100644
--- a/xarray/tests/test_plot.py
+++ b/xarray/tests/test_plot.py
@@ -2262,6 +2262,13 @@ class TestDatasetScatterPlots(PlotTestCase):
         with pytest.raises(ValueError):
             ds2.plot.scatter(x="A", y="B", hue="hue", hue_style="continuous")
 
+    def test_legend_labels(self):
+        # regression test for #4126: incorrect legend labels
+        ds2 = self.ds.copy()
+        ds2["hue"] = ["a", "a", "b", "b"]
+        lines = ds2.plot.scatter(x="A", y="B", hue="hue")
+        assert [t.get_text() for t in lines[0].axes.get_legend().texts] == ["a", "b"]
+
     def test_add_legend_by_default(self):
         sc = self.ds.plot.scatter(x="A", y="B", hue="hue")
         assert len(sc.figure.axes) == 2
diff --git a/xarray/tests/test_variable.py b/xarray/tests/test_variable.py
index 169105bb..efebe09e 100644
--- a/xarray/tests/test_variable.py
+++ b/xarray/tests/test_variable.py
@@ -1949,7 +1949,10 @@ class TestVariable(VariableSubclassobjects):
         # Test kept attrs
         with set_options(keep_attrs=True):
             new = Variable(["coord"], np.linspace(1, 10, 100), attrs=_attrs).coarsen(
-                windows={"coord": 1}, func=test_func, boundary="exact", side="left"
+                windows={"coord": 1},
+                func=test_func,
+                boundary="exact",
+                side="left",
             )
         assert new.attrs == _attrs