Add dot notation support (mdowds#24) (mdowds#47)

* Add dot notation support (mdowds#24)

* add readme doc

README.md

@@ -53,6 +53,7 @@ mock_db.collection('users').document('alovelace').set({
 })
 mock_db.collection('users').document('alovelace').set({'first': 'Augusta Ada'}, merge=True)
 mock_db.collection('users').document('alovelace').update({'born': 1815})
+mock_db.collection('users').document('alovelace').update({'favourite.color': 'red'})
 mock_db.collection('users').document('alovelace').update({'associates': ['Charles Babbage', 'Michael Faraday']})
 mock_db.collection('users').document('alovelace').collection('friends')
 mock_db.collection('users').document('alovelace').delete()

mockfirestore/_helpers.py

@@ -12,14 +12,23 @@
 Store = Dict[str, Collection]
 
 
-def get_by_path(data: Dict[str, T], path: Sequence[str]) -> T:
+def get_by_path(data: Dict[str, T], path: Sequence[str], create_nested: bool = False) -> T:
     """Access a nested object in root by item sequence."""
-    return reduce(operator.getitem, path, data)
 
+    def get_or_create(a, b):
+        if b not in a:
+            a[b] = {}
+        return a[b]
 
-def set_by_path(data: Dict[str, T], path: Sequence[str], value: T):
+    if create_nested:
+        return reduce(get_or_create, path, data)
+    else:
+        return reduce(operator.getitem, path, data)
+
+
+def set_by_path(data: Dict[str, T], path: Sequence[str], value: T, create_nested: bool = True):
     """Set a value in a nested object in root by item sequence."""
-    get_by_path(data, path[:-1])[path[-1]] = value
+    get_by_path(data, path[:-1], create_nested=True)[path[-1]] = value
 
 
 def delete_by_path(data: Dict[str, T], path: Sequence[str]):

mockfirestore/document.py

@@ -130,9 +130,15 @@ def _update_data(new_values: dict, default: Any):
             except (TypeError, KeyError):
                 item = default
 
-            set_by_path(data, path, item + value)
-    
+            set_by_path(data, path, item + value, create_nested=True)
+
     _update_data(increments, 0)
     _update_data(arr_unions, [])
 
-    document.update(data)
+    _apply_updates(document, data)
+
+
+def _apply_updates(document: Dict[str, Any], data: Dict[str, Any]):
+    for key, value in data.items():
+        path = key.split(".")
+        set_by_path(document, path, value, create_nested=True)

tests/test_document_reference.py

@@ -258,4 +258,42 @@ def test_document_update_transformerArrayUnionNonExistent(self):
         fs.collection('foo').document('first').update({'arr': firestore.ArrayUnion([1])})
 
         doc = fs.collection('foo').document('first').get().to_dict()
-        self.assertEqual(doc, {'arr': [1], 'spicy': 'tuna'})
+        self.assertEqual(doc, {'arr': [1], 'spicy': 'tuna'})
+
+    def test_document_update_nestedFieldDotNotation(self):
+        fs = MockFirestore()
+        fs._data = {"foo": {"first": {"nested": {"value": 1, "unchanged": "foo"}}}}
+
+        fs.collection("foo").document("first").update({"nested.value": 2})
+
+        doc = fs.collection("foo").document("first").get().to_dict()
+        self.assertEqual(doc, {"nested": {"value": 2, "unchanged": "foo"}})
+
+    def test_document_update_nestedFieldDotNotationNestedFieldCreation(self):
+        fs = MockFirestore()
+        fs._data = {"foo": {"first": {"other": None}}}  # non-existent nested field is created
+
+        fs.collection("foo").document("first").update({"nested.value": 2})
+
+        doc = fs.collection("foo").document("first").get().to_dict()
+        self.assertEqual(doc, {"nested": {"value": 2}, "other": None})
+
+    def test_document_update_nestedFieldDotNotationMultipleNested(self):
+        fs = MockFirestore()
+        fs._data = {"foo": {"first": {"other": None}}}
+
+        fs.collection("foo").document("first").update({"nested.subnested.value": 42})
+
+        doc = fs.collection("foo").document("first").get().to_dict()
+        self.assertEqual(doc, {"nested": {"subnested": {"value": 42}}, "other": None})
+
+    def test_document_update_nestedFieldDotNotationMultipleNestedWithTransformer(self):
+        fs = MockFirestore()
+        fs._data = {"foo": {"first": {"other": None}}}
+
+        fs.collection("foo").document("first").update(
+            {"nested.subnested.value": firestore.ArrayUnion([1, 3])}
+        )
+
+        doc = fs.collection("foo").document("first").get().to_dict()
+        self.assertEqual(doc, {"nested": {"subnested": {"value": [1, 3]}}, "other": None})