From 876efc354b106c096c2661dde8e0f9749a2f85cd Mon Sep 17 00:00:00 2001
From: Danny Hermes <daniel.j.hermes@gmail.com>
Date: Tue, 16 Feb 2016 07:33:01 -0800
Subject: [PATCH] Adding Bigtable PartialRowData.

These objects will be used when data is read from a table.
---
 gcloud/bigtable/row_data.py      | 48 +++++++++++++++++
 gcloud/bigtable/test_row_data.py | 90 ++++++++++++++++++++++++++++++++
 2 files changed, 138 insertions(+)

diff --git a/gcloud/bigtable/row_data.py b/gcloud/bigtable/row_data.py
index 3dc01b7776d59..bf8f4ac16696c 100644
--- a/gcloud/bigtable/row_data.py
+++ b/gcloud/bigtable/row_data.py
@@ -15,7 +15,10 @@
 """Container for Google Cloud Bigtable Cells and Streaming Row Contents."""
 
 
+import six
+
 from gcloud._helpers import _datetime_from_microseconds
+from gcloud._helpers import _to_bytes
 
 
 class Cell(object):
@@ -61,3 +64,48 @@ def __eq__(self, other):
 
     def __ne__(self, other):
         return not self.__eq__(other)
+
+
+class PartialRowData(object):
+    """Representation of partial row in a Google Cloud Bigtable Table.
+
+    These are expected to be updated directly from a
+    :class:`._generated.bigtable_service_messages_pb2.ReadRowsResponse`
+
+    :type row_key: bytes
+    :param row_key: The key for the row holding the (partial) data.
+    """
+
+    def __init__(self, row_key):
+        self._row_key = row_key
+        self._cells = {}
+        self._committed = False
+        self._chunks_encountered = False
+
+    def __eq__(self, other):
+        if not isinstance(other, self.__class__):
+            return False
+        return (other._row_key == self._row_key and
+                other._committed == self._committed and
+                other._chunks_encountered == self._chunks_encountered and
+                other._cells == self._cells)
+
+    def __ne__(self, other):
+        return not self.__eq__(other)
+
+    def to_dict(self):
+        """Convert the cells to a dictionary.
+
+        This is intended to be used with HappyBase, so the column family and
+        column qualiers are combined (with ``:``).
+
+        :rtype: dict
+        :returns: Dictionary containing all the data in the cells of this row.
+        """
+        result = {}
+        for column_family_id, columns in six.iteritems(self._cells):
+            for column_qual, cells in six.iteritems(columns):
+                key = (_to_bytes(column_family_id) + b':' +
+                       _to_bytes(column_qual))
+                result[key] = cells
+        return result
diff --git a/gcloud/bigtable/test_row_data.py b/gcloud/bigtable/test_row_data.py
index 935bfb2b31d88..c05be8830c643 100644
--- a/gcloud/bigtable/test_row_data.py
+++ b/gcloud/bigtable/test_row_data.py
@@ -89,3 +89,93 @@ def test___ne__(self):
         cell1 = self._makeOne(value1, timestamp)
         cell2 = self._makeOne(value2, timestamp)
         self.assertNotEqual(cell1, cell2)
+
+
+class TestPartialRowData(unittest2.TestCase):
+
+    def _getTargetClass(self):
+        from gcloud.bigtable.row_data import PartialRowData
+        return PartialRowData
+
+    def _makeOne(self, *args, **kwargs):
+        return self._getTargetClass()(*args, **kwargs)
+
+    def test_constructor(self):
+        row_key = object()
+        partial_row_data = self._makeOne(row_key)
+        self.assertTrue(partial_row_data._row_key is row_key)
+        self.assertEqual(partial_row_data._cells, {})
+        self.assertFalse(partial_row_data._committed)
+        self.assertFalse(partial_row_data._chunks_encountered)
+
+    def test___eq__(self):
+        row_key = object()
+        partial_row_data1 = self._makeOne(row_key)
+        partial_row_data2 = self._makeOne(row_key)
+        self.assertEqual(partial_row_data1, partial_row_data2)
+
+    def test___eq__type_differ(self):
+        partial_row_data1 = self._makeOne(None)
+        partial_row_data2 = object()
+        self.assertNotEqual(partial_row_data1, partial_row_data2)
+
+    def test___ne__same_value(self):
+        row_key = object()
+        partial_row_data1 = self._makeOne(row_key)
+        partial_row_data2 = self._makeOne(row_key)
+        comparison_val = (partial_row_data1 != partial_row_data2)
+        self.assertFalse(comparison_val)
+
+    def test___ne__(self):
+        row_key1 = object()
+        partial_row_data1 = self._makeOne(row_key1)
+        row_key2 = object()
+        partial_row_data2 = self._makeOne(row_key2)
+        self.assertNotEqual(partial_row_data1, partial_row_data2)
+
+    def test___ne__committed(self):
+        row_key = object()
+        partial_row_data1 = self._makeOne(row_key)
+        partial_row_data1._committed = object()
+        partial_row_data2 = self._makeOne(row_key)
+        self.assertNotEqual(partial_row_data1, partial_row_data2)
+
+    def test___ne__cells(self):
+        row_key = object()
+        partial_row_data1 = self._makeOne(row_key)
+        partial_row_data1._cells = object()
+        partial_row_data2 = self._makeOne(row_key)
+        self.assertNotEqual(partial_row_data1, partial_row_data2)
+
+    def test_to_dict(self):
+        cell1 = object()
+        cell2 = object()
+        cell3 = object()
+
+        family_name1 = u'name1'
+        family_name2 = u'name2'
+        qual1 = b'col1'
+        qual2 = b'col2'
+        qual3 = b'col3'
+
+        partial_row_data = self._makeOne(None)
+        partial_row_data._cells = {
+            family_name1: {
+                qual1: cell1,
+                qual2: cell2,
+            },
+            family_name2: {
+                qual3: cell3,
+            },
+        }
+
+        result = partial_row_data.to_dict()
+        col1 = family_name1.encode('ascii') + b':' + qual1
+        col2 = family_name1.encode('ascii') + b':' + qual2
+        col3 = family_name2.encode('ascii') + b':' + qual3
+        expected_result = {
+            col1: cell1,
+            col2: cell2,
+            col3: cell3,
+        }
+        self.assertEqual(result, expected_result)