Merge pull request #1490 from dhermes/bigtable-partial-row-props

dhermes · dhermes · commit d3328098ab8e · 2016-02-18T11:36:19.000-08:00
Adding clear() to Bigtable row data.
diff --git a/gcloud/bigtable/row_data.py b/gcloud/bigtable/row_data.py
@@ -15,6 +15,7 @@
 """Container for Google Cloud Bigtable Cells and Streaming Row Contents."""
 
 
+import copy
 import six
 
 from gcloud._helpers import _datetime_from_microseconds
@@ -110,6 +111,42 @@ def to_dict(self):
                 result[key] = cells
         return result
 
+    @property
+    def cells(self):
+        """Property returning all the cells accumulated on this partial row.
+
+        :rtype: dict
+        :returns: Dictionary of the :class:`Cell` objects accumulated. This
+                  dictionary has two-levels of keys (first for column families
+                  and second for column names/qualifiers within a family). For
+                  a given column, a list of :class:`Cell` objects is stored.
+        """
+        return copy.deepcopy(self._cells)
+
+    @property
+    def row_key(self):
+        """Getter for the current (partial) row's key.
+
+        :rtype: bytes
+        :returns: The current (partial) row's key.
+        """
+        return self._row_key
+
+    @property
+    def committed(self):
+        """Getter for the committed status of the (partial) row.
+
+        :rtype: bool
+        :returns: The committed status of the (partial) row.
+        """
+        return self._committed
+
+    def clear(self):
+        """Clears all cells that have been added."""
+        self._committed = False
+        self._chunks_encountered = False
+        self._cells.clear()
+
 
 class PartialRowsData(object):
     """Convenience wrapper for consuming a ``ReadRows`` streaming response.
diff --git a/gcloud/bigtable/test_row_data.py b/gcloud/bigtable/test_row_data.py
@@ -177,6 +177,36 @@ def test_to_dict(self):
         }
         self.assertEqual(result, expected_result)
 
+    def test_cells_property(self):
+        partial_row_data = self._makeOne(None)
+        cells = {1: 2}
+        partial_row_data._cells = cells
+        # Make sure we get a copy, not the original.
+        self.assertFalse(partial_row_data.cells is cells)
+        self.assertEqual(partial_row_data.cells, cells)
+
+    def test_row_key_getter(self):
+        row_key = object()
+        partial_row_data = self._makeOne(row_key)
+        self.assertTrue(partial_row_data.row_key is row_key)
+
+    def test_committed_getter(self):
+        partial_row_data = self._makeOne(None)
+        partial_row_data._committed = value = object()
+        self.assertTrue(partial_row_data.committed is value)
+
+    def test_clear(self):
+        partial_row_data = self._makeOne(None)
+        cells = {1: 2}
+        partial_row_data._cells = cells
+        self.assertEqual(partial_row_data.cells, cells)
+        partial_row_data._committed = True
+        partial_row_data._chunks_encountered = True
+        partial_row_data.clear()
+        self.assertFalse(partial_row_data.committed)
+        self.assertFalse(partial_row_data._chunks_encountered)
+        self.assertEqual(partial_row_data.cells, {})
+
 
 class TestPartialRowsData(unittest2.TestCase):
 
