Remove extract_field_paths

firestore/google/cloud/firestore_v1beta1/_helpers.py

@@ -819,54 +819,35 @@ def process_server_timestamp(document_data, split_on_dots=True):
         * The remaining keys in ``document_data`` after removing the
           server timestamp sentinels
     """
+    field_paths = []
     transform_paths = []
     actual_data = {}
     for field_name, value in six.iteritems(document_data):
         if isinstance(value, dict):
-            sub_field_paths, sub_data = process_server_timestamp(value, False)
-            for sub_path in sub_field_paths:
+            sub_transform_paths, sub_data, sub_field_paths = process_server_timestamp(value, False)
+            for sub_path in sub_transform_paths:
                 field_path = FieldPath.from_string(field_name)
                 field_path.parts = field_path.parts + sub_path.parts
                 transform_paths.extend([field_path])
             if sub_data:
                 # Only add a key to ``actual_data`` if there is data.
+
                 actual_data[field_name] = sub_data
+                for sub_field_path in sub_field_paths:
+                    field_path = FieldPath.from_string(field_name)
+                    field_path.parts = field_path.parts + sub_field_path.parts
+                    field_paths.append(field_path)
         elif value is constants.SERVER_TIMESTAMP:
             if split_on_dots:
                 transform_paths.append(FieldPath(*field_name.split(".")))
             else:
                 transform_paths.append(FieldPath.from_string(field_name))
         else:
             actual_data[field_name] = value
+            field_paths.append(FieldPath(field_name))
     if not transform_paths:
         actual_data = document_data
-    return transform_paths, actual_data
-
-
-def extract_field_paths(document_data):
-    """Extract field paths from document data
-
-    Args:
-       document_data (dict): The dictionary of the actual set data.
-
-    Returns:
-       List[~.firestore_v1beta1._helpers.FieldPath]:
-           A list of `FieldPath` instances from the actual data.
-    """
-    field_paths = []
-    for field_name, value in six.iteritems(document_data):
-        if isinstance(value, dict):
-            sub_field_paths = extract_field_paths(value)
-            for sub_path in sub_field_paths:
-                paths = [field_name]
-                paths.extend(sub_path.parts)
-                field_path = FieldPath(*paths)
-                field_paths.append(field_path)
-        else:
-            path = FieldPath(field_name)
-            field_paths.append(path)
-    return field_paths
-
+    return transform_paths, actual_data, field_paths
 
 
 def get_transform_pb(document_path, transform_paths):
@@ -913,7 +894,7 @@ def pbs_for_set(document_path, document_data, merge=False, exists=None):
         List[google.cloud.firestore_v1beta1.types.Write]: One
         or two ``Write`` protobuf instances for ``set()``.
     """
-    transform_paths, actual_data = process_server_timestamp(
+    transform_paths, actual_data, field_paths = process_server_timestamp(
         document_data, False)
     update_pb = write_pb2.Write(
         update=document_pb2.Document(
@@ -926,7 +907,6 @@ def pbs_for_set(document_path, document_data, merge=False, exists=None):
             common_pb2.Precondition(exists=exists))
 
     if merge:
-        field_paths = extract_field_paths(document_data)
         field_paths = canonicalize_field_paths(field_paths)
         mask = common_pb2.DocumentMask(field_paths=sorted(field_paths))
         update_pb.update_mask.CopyFrom(mask)
@@ -985,7 +965,7 @@ def pbs_for_update(client, document_path, field_updates, option):
         # Default uses ``exists=True``.
         option = client.write_option(exists=True)
 
-    transform_paths, actual_updates = process_server_timestamp(field_updates)
+    transform_paths, actual_updates, field_paths = process_server_timestamp(field_updates)
     if not (transform_paths or actual_updates):
         raise ValueError('There are only ServerTimeStamp objects or is empty.')
     update_values, field_paths = FieldPathHelper.to_field_paths(actual_updates)

firestore/tests/unit/test__helpers.py

@@ -1284,14 +1284,21 @@ def _call_fut(document_data):
 
     def test_no_fields(self):
         import collections
+        from google.cloud.firestore_v1beta1 import _helpers
 
         data = collections.OrderedDict((
             ('one', 1),
             ('two', 2.25),
             ('three', [False, True, True]),
         ))
-        field_paths, actual_data = self._call_fut(data)
-        self.assertEqual(field_paths, [])
+        expected_field_paths = [
+            _helpers.FieldPath('one'),
+            _helpers.FieldPath('two'),
+            _helpers.FieldPath('three')
+        ]
+        transform_paths, actual_data, field_paths = self._call_fut(data)
+        self.assertEqual(transform_paths, [])
+        self.assertEqual(field_paths, expected_field_paths)
         self.assertIs(actual_data, data)
 
     def test_simple_fields(self):
@@ -1313,17 +1320,26 @@ def test_simple_fields(self):
             ('top5', 200),
             ('top6', nested2),
         ))
-        field_paths, actual_data = self._call_fut(data)
-        self.assertEqual(
-            field_paths, [_helpers.FieldPath('top1', 'bottom2'),
-                          _helpers.FieldPath('top4'),
-                          _helpers.FieldPath('top6', 'bottom7')])
+        expected_transform_paths = [
+            _helpers.FieldPath('top1', 'bottom2'),
+            _helpers.FieldPath('top4'),
+            _helpers.FieldPath('top6', 'bottom7')
+        ]
+        expected_field_paths = [
+            _helpers.FieldPath('top1', 'bottom3'),
+            _helpers.FieldPath('top5')]
         expected_data = {
             'top1': {
                 'bottom3': data['top1']['bottom3'],
             },
             'top5': data['top5'],
         }
+        transform_paths, actual_data, field_paths = self._call_fut(data)
+        self.assertEqual(
+            transform_paths,
+            expected_transform_paths
+        )
+        self.assertEqual(field_paths, expected_field_paths)
         self.assertEqual(actual_data, expected_data)
 
     def test_field_updates(self):
@@ -1337,9 +1353,9 @@ def test_field_updates(self):
             ('c.d', {'e': SERVER_TIMESTAMP}),
             ('f.g', SERVER_TIMESTAMP),
         ))
-        field_paths, actual_data = self._call_fut(data)
-        self.assertEqual(field_paths, [_helpers.FieldPath('c', 'd', 'e'),
-                                       _helpers.FieldPath('f', 'g')])
+        transform_paths, actual_data, field_paths = self._call_fut(data)
+        self.assertEqual(transform_paths, [_helpers.FieldPath('c', 'd', 'e'),
+                                           _helpers.FieldPath('f', 'g')])
 
         expected_data = {'a': {'b': data['a']['b']}}
         self.assertEqual(actual_data, expected_data)

firestore/tests/unit/test_document.py

@@ -277,7 +277,7 @@ def _write_pb_for_set(document_path, document_data, merge):
             ),
         )
         if merge:
-            field_paths = _helpers.extract_field_paths(document_data)
+            _, _, field_paths = _helpers.process_server_timestamp(document_data)
             field_paths = _helpers.canonicalize_field_paths(field_paths)
             mask = common_pb2.DocumentMask(field_paths=sorted(field_paths))
             write_pbs.update_mask.CopyFrom(mask)