Add async support and format with black

README.md

@@ -1,6 +1,6 @@
 # Python Mock Firestore
 
-An in-memory implementation of the [Python client library](https://github.com/googleapis/python-firestore) for Google Cloud Firestore, intended for use in tests to replace the real thing. This project is in early stages and is only a partial implementation of the real  client library.
+An in-memory implementation of the [Python client library](https://github.com/googleapis/python-firestore) for Google Cloud Firestore, intended for use in tests to replace the real thing. This project is only a partial implementation of the real client library.
 
 To install:
 
@@ -10,6 +10,8 @@ Python 3.6+ is required for it to work.
 
 ## Usage
 
+### Sync
+
 ```python
 db = firestore.Client()
 mock_db = MockFirestore()
@@ -19,14 +21,28 @@ db.collection('users').get()
 mock_db.collection('users').get()
 ```
 
+### Async
+
+```python
+db = firestore.AsyncClient()
+mock_db = AsyncMockFirestore()
+
+await db.collection('users').get()
+await mock_db.collection('users').get()
+```
+
 To reset the store to an empty state, use the `reset()` method:
 ```python
 mock_db = MockFirestore()
 mock_db.reset()
 ```
 
+or the equivalent method of `AsyncMockFirestore`
+
 ## Supported operations
 
+### Sync
+
 ```python
 mock_db = MockFirestore()
 
@@ -57,7 +73,7 @@ mock_db.collection('users').document('alovelace').update({'favourite.color': 're
 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()
-mock_db.collection('users').document(document_id: 'alovelace').delete()
+mock_db.collection('users').document('alovelace').delete()
 mock_db.collection('users').add({'first': 'Ada', 'last': 'Lovelace'}, 'alovelace')
 mock_db.get_all([mock_db.collection('users').document('alovelace')])
 mock_db.document('users/alovelace')
@@ -104,6 +120,81 @@ transaction.delete(mock_db.collection('users').document('alovelace'))
 transaction.commit()
 ```
 
+### Async
+*(Where usage of those differs from the above)*
+
+*Note: all iterator methods like `stream` or `list_documents` in AsyncMockFirestore and its associated async classes 
+return asynchronous iterators, so when iterating over them,
+`async for` syntax must be used.*
+
+```python
+mock_db = AsyncMockFirestore()
+
+# Collections
+await mock_db.collection('users').get()
+
+# async iterators
+[doc_ref async for doc_ref in mock_db.collection('users').list_documents()]
+[doc_snapshot async for doc_snapshot in mock_db.collection('users').stream()]
+
+# Documents
+await mock_db.collection('users').document('alovelace').get()
+doc_snapshot = await mock_db.collection('users').document('alovelace').get()
+doc_snapshot.exists
+doc_snapshot.to_dict()
+await mock_db.collection('users').document('alovelace').set({
+    'first': 'Ada',
+    'last': 'Lovelace'
+})
+await mock_db.collection('users').document('alovelace').set({'first': 'Augusta Ada'}, merge=True)
+await mock_db.collection('users').document('alovelace').update({'born': 1815})
+await mock_db.collection('users').document('alovelace').update({'favourite.color': 'red'})
+await mock_db.collection('users').document('alovelace').update({'associates': ['Charles Babbage', 'Michael Faraday']})
+await mock_db.collection('users').document('alovelace').delete()
+await mock_db.collection('users').document('alovelace').delete()
+await mock_db.collection('users').add({'first': 'Ada', 'last': 'Lovelace'}, 'alovelace')
+await mock_db.get_all([mock_db.collection('users').document('alovelace')])
+await mock_db.document('users/alovelace').update({'born': 1815})
+
+# Querying
+await mock_db.collection('users').order_by('born').get()
+await mock_db.collection('users').order_by('born', direction='DESCENDING').get()
+await mock_db.collection('users').limit(5).get()
+await mock_db.collection('users').where('born', '==', 1815).get()
+await mock_db.collection('users').where('born', '!=', 1815).get()
+await mock_db.collection('users').where('born', '<', 1815).get()
+await mock_db.collection('users').where('born', '>', 1815).get()
+await mock_db.collection('users').where('born', '<=', 1815).get()
+await mock_db.collection('users').where('born', '>=', 1815).get()
+
+# async iterators
+mock_db.collection('users').where('born', 'in', [1815, 1900]).stream()
+mock_db.collection('users').where('born', 'in', [1815, 1900]).stream()
+mock_db.collection('users').where('associates', 'array_contains', 'Charles Babbage').stream()
+mock_db.collection('users').where('associates', 'array_contains_any', ['Charles Babbage', 'Michael Faraday']).stream()
+
+# Transforms
+await mock_db.collection('users').document('alovelace').update({'likes': firestore.Increment(1)})
+await mock_db.collection('users').document('alovelace').update({'associates': firestore.ArrayUnion(['Andrew Cross', 'Charles Wheatstone'])})
+await mock_db.collection('users').document('alovelace').update({firestore.DELETE_FIELD: "born"})
+await mock_db.collection('users').document('alovelace').update({'associates': firestore.ArrayRemove(['Andrew Cross'])})
+
+
+# Transactions
+transaction = mock_db.transaction()
+transaction.id
+transaction.in_progress
+await transaction.get(mock_db.collection('users').where('born', '==', 1815))
+await transaction.get(mock_db.collection('users').document('alovelace'))
+await transaction.get_all([mock_db.collection('users').document('alovelace')])
+
+transaction.set(mock_db.collection('users').document('alovelace'), {'born': 1815})
+transaction.update(mock_db.collection('users').document('alovelace'), {'born': 1815})
+transaction.delete(mock_db.collection('users').document('alovelace'))
+await transaction.commit()
+```
+
+
 ## Running the tests
 * Create and activate a virtualenv with a Python version of at least 3.6
 * Install dependencies with `pip install -r requirements-dev-minimal.txt`
@@ -113,7 +204,7 @@ transaction.commit()
 
 * [Matt Dowds](https://github.com/mdowds)
 * [Chris Tippett](https://github.com/christippett)
-* [Anton Melnikov](https://github.com/notnami)
+* [Anna Melnikov](https://github.com/anna-hope)
 * [Ben Riggleman](https://github.com/briggleman)
 * [Steve Atwell](https://github.com/satwell)
 * [ahti123](https://github.com/ahti123)

mockfirestore/__init__.py

@@ -13,3 +13,9 @@
 from mockfirestore.query import Query
 from mockfirestore._helpers import Timestamp
 from mockfirestore.transaction import Transaction
+
+from mockfirestore.async_client import AsyncMockFirestore
+from mockfirestore.async_document import AsyncDocumentReference
+from mockfirestore.async_collection import AsyncCollectionReference
+from mockfirestore.async_query import AsyncQuery
+from mockfirestore.async_transaction import AsyncTransaction

mockfirestore/_helpers.py

@@ -3,16 +3,18 @@
 import string
 from datetime import datetime as dt
 from functools import reduce
-from typing import (Dict, Any, Tuple, TypeVar, Sequence, Iterator)
+from typing import Dict, Any, Tuple, TypeVar, Sequence, Iterator, AsyncIterable, List
 
-T = TypeVar('T')
+T = TypeVar("T")
 KeyValuePair = Tuple[str, Dict[str, Any]]
 Document = Dict[str, Any]
 Collection = Dict[str, Document]
 Store = Dict[str, Collection]
 
 
-def get_by_path(data: Dict[str, T], path: Sequence[str], create_nested: bool = False) -> 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."""
 
     def get_or_create(a, b):
@@ -26,7 +28,9 @@ def get_or_create(a, b):
         return reduce(operator.getitem, path, data)
 
 
-def set_by_path(data: Dict[str, T], path: Sequence[str], value: T, create_nested: bool = True):
+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], create_nested=True)[path[-1]] = value
 
@@ -37,7 +41,9 @@ def delete_by_path(data: Dict[str, T], path: Sequence[str]):
 
 
 def generate_random_string():
-    return ''.join(random.choice(string.ascii_letters + string.digits) for _ in range(20))
+    return "".join(
+        random.choice(string.ascii_letters + string.digits) for _ in range(20)
+    )
 
 
 class Timestamp:
@@ -55,14 +61,16 @@ def from_now(cls):
 
     @property
     def seconds(self):
-        return str(self._timestamp).split('.')[0]
+        return str(self._timestamp).split(".")[0]
 
     @property
     def nanos(self):
-        return str(self._timestamp).split('.')[1]
+        return str(self._timestamp).split(".")[1]
 
 
-def get_document_iterator(document: Dict[str, Any], prefix: str = '') -> Iterator[Tuple[str, Any]]:
+def get_document_iterator(
+    document: Dict[str, Any], prefix: str = ""
+) -> Iterator[Tuple[str, Any]]:
     """
     :returns: (dot-delimited path, value,)
     """
@@ -74,4 +82,8 @@ def get_document_iterator(document: Dict[str, Any], prefix: str = '') -> Iterato
         if not prefix:
             yield key, value
         else:
-            yield '{}.{}'.format(prefix, key), value
+            yield "{}.{}".format(prefix, key), value
+
+
+async def consume_async_iterable(iterable: AsyncIterable[T]) -> List[T]:
+    return [item async for item in iterable]

mockfirestore/async_client.py

@@ -0,0 +1,45 @@
+from typing import AsyncIterable, Iterable
+
+from mockfirestore.async_document import AsyncDocumentReference
+from mockfirestore.async_collection import AsyncCollectionReference
+from mockfirestore.async_transaction import AsyncTransaction
+from mockfirestore.client import MockFirestore
+from mockfirestore.document import DocumentSnapshot
+
+
+class AsyncMockFirestore(MockFirestore):
+    def document(self, path: str) -> AsyncDocumentReference:
+        doc = super().document(path)
+        assert isinstance(doc, AsyncDocumentReference)
+        return doc
+
+    def collection(self, path: str) -> AsyncCollectionReference:
+        path = path.split("/")
+
+        if len(path) % 2 != 1:
+            raise Exception("Cannot create collection at path {}".format(path))
+
+        name = path[-1]
+        if len(path) > 1:
+            current_position = self._ensure_path(path)
+            return current_position.collection(name)
+        else:
+            if name not in self._data:
+                self._data[name] = {}
+            return AsyncCollectionReference(self._data, [name])
+
+    async def collections(self) -> AsyncIterable[AsyncCollectionReference]:
+        for collection_name in self._data:
+            yield AsyncCollectionReference(self._data, [collection_name])
+
+    async def get_all(
+        self,
+        references: Iterable[AsyncDocumentReference],
+        field_paths=None,
+        transaction=None,
+    ) -> AsyncIterable[DocumentSnapshot]:
+        for doc_ref in set(references):
+            yield await doc_ref.get()
+
+    def transaction(self, **kwargs) -> AsyncTransaction:
+        return AsyncTransaction(self, **kwargs)

mockfirestore/async_collection.py

@@ -0,0 +1,70 @@
+from typing import Optional, List, Tuple, Dict, AsyncIterator, Any, Union
+from mockfirestore.async_document import AsyncDocumentReference
+from mockfirestore.async_query import AsyncQuery
+from mockfirestore.collection import CollectionReference
+from mockfirestore.document import DocumentSnapshot, DocumentReference
+from mockfirestore._helpers import Timestamp, get_by_path
+
+
+class AsyncCollectionReference(CollectionReference):
+    def document(self, document_id: Optional[str] = None) -> AsyncDocumentReference:
+        doc_ref = super().document(document_id)
+        return AsyncDocumentReference(
+            doc_ref._data, doc_ref._path, parent=doc_ref.parent
+        )
+
+    async def get(self) -> List[DocumentSnapshot]:
+        return super().get()
+
+    async def add(
+        self, document_data: Dict, document_id: str = None
+    ) -> Tuple[Timestamp, AsyncDocumentReference]:
+        timestamp, doc_ref = super().add(document_data, document_id=document_id)
+        async_doc_ref = AsyncDocumentReference(
+            doc_ref._data, doc_ref._path, parent=doc_ref.parent
+        )
+        return timestamp, async_doc_ref
+
+    async def list_documents(
+        self, page_size: Optional[int] = None
+    ) -> AsyncIterator[DocumentReference]:
+        docs = super().list_documents()
+        for doc in docs:
+            yield doc
+
+    async def stream(self, transaction=None) -> AsyncIterator[DocumentSnapshot]:
+        for key in sorted(get_by_path(self._data, self._path)):
+            doc_snapshot = await self.document(key).get()
+            yield doc_snapshot
+
+    def where(self, field: str, op: str, value: Any) -> AsyncQuery:
+        query = AsyncQuery(self, field_filters=[(field, op, value)])
+        return query
+
+    def order_by(self, key: str, direction: Optional[str] = None) -> AsyncQuery:
+        query = AsyncQuery(self, orders=[(key, direction)])
+        return query
+
+    def limit(self, limit_amount: int) -> AsyncQuery:
+        query = AsyncQuery(self, limit=limit_amount)
+        return query
+
+    def offset(self, offset: int) -> AsyncQuery:
+        query = AsyncQuery(self, offset=offset)
+        return query
+
+    def start_at(self, document_fields_or_snapshot: Union[dict, DocumentSnapshot]) -> AsyncQuery:
+        query = AsyncQuery(self, start_at=(document_fields_or_snapshot, True))
+        return query
+
+    def start_after(self, document_fields_or_snapshot: Union[dict, DocumentSnapshot]) -> AsyncQuery:
+        query = AsyncQuery(self, start_at=(document_fields_or_snapshot, False))
+        return query
+
+    def end_at(self, document_fields_or_snapshot: Union[dict, DocumentSnapshot]) -> AsyncQuery:
+        query = AsyncQuery(self, end_at=(document_fields_or_snapshot, True))
+        return query
+
+    def end_before(self, document_fields_or_snapshot: Union[dict, DocumentSnapshot]) -> AsyncQuery:
+        query = AsyncQuery(self, end_at=(document_fields_or_snapshot, False))
+        return query

mockfirestore/async_document.py

@@ -0,0 +1,29 @@
+from copy import deepcopy
+from typing import Dict, Any
+from mockfirestore import NotFound
+from mockfirestore.document import DocumentReference, DocumentSnapshot
+
+
+class AsyncDocumentReference(DocumentReference):
+    async def get(self) -> DocumentSnapshot:
+        return super().get()
+
+    async def delete(self):
+        super().delete()
+
+    async def set(self, data: Dict[str, Any], merge=False):
+        if merge:
+            try:
+                await self.update(deepcopy(data))
+            except NotFound:
+                await self.set(data)
+        else:
+            super().set(data, merge=merge)
+
+    async def update(self, data: Dict[str, Any]):
+        super().update(data)
+
+    def collection(self, name) -> 'AsyncCollectionReference':
+        from mockfirestore.async_collection import AsyncCollectionReference
+        coll_ref = super().collection(name)
+        return AsyncCollectionReference(coll_ref._data, coll_ref._path, self)

mockfirestore/async_query.py

@@ -0,0 +1,16 @@
+from typing import List, AsyncIterator
+from mockfirestore.document import DocumentSnapshot
+from mockfirestore.query import Query
+from mockfirestore._helpers import consume_async_iterable
+
+
+class AsyncQuery(Query):
+    async def stream(self, transaction=None) -> AsyncIterator[DocumentSnapshot]:
+        doc_snapshots = await consume_async_iterable(self.parent.stream())
+        doc_snapshots = super()._process_field_filters(doc_snapshots)
+        doc_snapshots = super()._process_pagination(doc_snapshots)
+        for doc_snapshot in doc_snapshots:
+            yield doc_snapshot
+
+    async def get(self, transaction=None) -> List[DocumentSnapshot]:
+        return await consume_async_iterable(self.stream())