dask: Data.flat

cf/data/data.py

@@ -10135,40 +10135,43 @@ def last_element(self):
             "setting the log level to 'DEBUG'."
         )
 
+    @daskified(_DASKIFIED_VERBOSE)
     def flat(self, ignore_masked=True):
         """Return a flat iterator over elements of the data array.
 
+        **Performance**
+
+        Any delayed operations and/or disk interactions will be
+        executed during *each* iteration, possibly leading to poor
+        performance. If possible, consider bringing the values into
+        memory first with `persist` or using ``d.array.flat``.
+
+        .. seealso:: `flatten`, `persist`
+
         :Parameters:
 
             ignore_masked: `bool`, optional
                 If False then masked and unmasked elements will be
-                returned. By default only unmasked elements are returned
+                returned. By default only unmasked elements are
+                returned
 
         :Returns:
 
             generator
                 An iterator over elements of the data array.
 
-        **Examples:**
+        **Examples**
 
+        >>> d = cf.Data([[1, 2], [3,4]], mask=[[0, 1], [0, 0]])
         >>> print(d.array)
-        [[1 -- 3]]
-        >>> for x in d.flat():
-        ...     print(x)
-        ...
-        1
-        3
-
-        >>> for x in d.flat(ignore_masked=False):
-        ...     print(x)
-        ...
-        1
-        --
-        3
+        [[1 --]
+         [3 4]]
+        >>> list(d.flat())
+        [1, 3, 4]
+        >>> list(d.flat(ignore_masked=False))
+        [1, masked, 3, 4]
 
         """
-        self.to_memory()
-
         mask = self.mask
 
         if ignore_masked:

cf/test/test_Data.py

@@ -3956,6 +3956,13 @@ def test_Data_set_units(self):
         with self.assertRaises(ValueError):
             d.set_units("km")
 
+    def test_Data_flat(self):
+        d = cf.Data([[1, 2], [3, 4]], mask=[[0, 1], [0, 0]])
+        self.assertEqual(list(d.flat()), [1, 3, 4])
+        self.assertEqual(
+            list(d.flat(ignore_masked=False)), [1, np.ma.masked, 3, 4]
+        )
+
     @unittest.skipIf(TEST_DASKIFIED_ONLY, "Needs updated NetCDFArray to test")
     def test_Data_get_filenames(self):
         pass
@@ -3972,7 +3979,7 @@ def test_Data_data(self):
             cf.Data(1),
             cf.Data([1, 2], fill_value=0),
             cf.Data([1, 2], "m"),
-            cf.Data([1, 2], mask=[1,0], units="m"),
+            cf.Data([1, 2], mask=[1, 0], units="m"),
             cf.Data([[0, 1, 2], [3, 4, 5]], chunks=2),
         ]:
             self.assertIs(d.data, d)