From 6b8c1b0b02866775675358aa6b2c78d3f6593ca7 Mon Sep 17 00:00:00 2001
From: 0x0L <0x0L@github.com>
Date: Sun, 10 Dec 2017 15:37:23 +0100
Subject: [PATCH] dataset fix and more tests

---
 xarray/core/dataset.py        | 10 +++++-----
 xarray/tests/test_dataset.py  | 12 +++++++++++-
 xarray/tests/test_variable.py |  7 +++++++
 3 files changed, 23 insertions(+), 6 deletions(-)

diff --git a/xarray/core/dataset.py b/xarray/core/dataset.py
index aad99be6ab2..e8eb71c7ed1 100644
--- a/xarray/core/dataset.py
+++ b/xarray/core/dataset.py
@@ -3278,11 +3278,11 @@ def rank(self, dim, pct=False, keep_attrs=False):
 
         variables = OrderedDict()
         for name, var in iteritems(self.variables):
-            if name in self.data_vars and dim in var.dims:
-                variables[name] = var.rank(dim, pct=pct)
-                variables.update({
-                    k: self.variables[k] for k in var.dims
-                    if k not in variables and k in self.variables})
+            if name in self.data_vars:
+                if dim in var.dims:
+                    variables[name] = var.rank(dim, pct=pct)
+            else:
+                variables[name] = var
 
         coord_names = set(k for k in self.coords if k in variables)
         attrs = self.attrs if keep_attrs else None
diff --git a/xarray/tests/test_dataset.py b/xarray/tests/test_dataset.py
index 588a98f1437..b62d02127e6 100644
--- a/xarray/tests/test_dataset.py
+++ b/xarray/tests/test_dataset.py
@@ -3417,9 +3417,19 @@ def test_quantile(self):
     @requires_bottleneck
     def test_rank(self):
         ds = create_test_data(seed=1234)
-        x = ds.rank('dim3').var3
+        # only ds.var3 depends on dim3
+        z = ds.rank('dim3')
+        self.assertItemsEqual(['var3'], list(z.data_vars))
+        # same as dataarray version
+        x = z.var3
         y = ds.var3.rank('dim3')
         self.assertDataArrayEqual(x, y)
+        # coordinates stick
+        self.assertItemsEqual(list(z.coords), list(ds.coords))
+        self.assertItemsEqual(list(x.coords), list(y.coords))
+        # invalid dim
+        with raises_regex(ValueError, 'does not contain'):
+            x.rank('invalid_dim')
 
     def test_count(self):
         ds = Dataset({'x': ('a', [np.nan, 1]), 'y': 0, 'z': np.nan})
diff --git a/xarray/tests/test_variable.py b/xarray/tests/test_variable.py
index eacf98b399f..3d4a8eb2c5b 100644
--- a/xarray/tests/test_variable.py
+++ b/xarray/tests/test_variable.py
@@ -1376,6 +1376,13 @@ def test_rank(self):
         v = Variable(['x'], [3.0, 1.0, np.nan, 2.0, 4.0])
         v_expect = Variable(['x'], [0.75, 0.25, np.nan, 0.5, 1.0])
         self.assertVariableEqual(v.rank('x', pct=True), v_expect)
+        # no dask arrays
+        w = v.chunk(2)
+        with raises_regex(TypeError, 'arrays stored as dask'):
+            w.rank('x')
+        # invalid dim
+        with raises_regex(ValueError, 'not found'):
+            v.rank('y')
 
     def test_big_endian_reduce(self):
         # regression test for GH489