From a73c8fa9fca8b7c5edae342c955aba5a68536891 Mon Sep 17 00:00:00 2001 From: Makoto Tanji Date: Fri, 16 Aug 2024 15:13:21 +0900 Subject: [PATCH 1/2] Add all, to_a --- pyfireconsole/models/pyfire_model.py | 29 ++++++++++++ pyfireconsole/queries/query_runner.py | 16 ++++++- tests/test_usecase.py | 65 +++++++++++++++++++++++++++ 3 files changed, 109 insertions(+), 1 deletion(-) diff --git a/pyfireconsole/models/pyfire_model.py b/pyfireconsole/models/pyfire_model.py index f1d87be..e3aabc3 100644 --- a/pyfireconsole/models/pyfire_model.py +++ b/pyfireconsole/models/pyfire_model.py @@ -18,6 +18,7 @@ class PyfireCollection(Generic[ModelType]): _collection: Optional[Iterable[dict]] = None _where_cond: Optional[WhereCondition] = None _order_cond: Optional[OrderCondition] = None + _limit: int = 1000 # default limit to prevent loading too large collections def __init__(self, model_class: Type[ModelType]): self.model_class = model_class @@ -61,6 +62,9 @@ def __iter__(self): if self._order_cond is not None: query = query.order(self._order_cond.field, self._order_cond.direction) + + self._collection = query.limit(self._limit).iter() + self._collection = query.iter() for doc in self._collection: @@ -132,6 +136,31 @@ def order(self, field: str, direction: str = 'ASCENDING') -> 'PyfireCollection[M return coll + def all(self) -> 'PyfireCollection[ModelType]': + """ + Get all documents in the collection. + + Returns: + PyfireCollection[ModelType]: A new PyfireCollection instance containing all documents. + """ + coll = PyfireCollection(self.model_class) + coll._where_cond = self._where_cond + coll._order_cond = self._order_cond + coll._limit = self._limit + if self._parent_model is not None: + coll.set_parent(self._parent_model) + + return coll + + def to_a(self, limit: int = 1000) -> list[ModelType]: + """ + Convert the collection to a list. + + Returns: + list[ModelType]: The collection as a list. + """ + return list(self) + def add(self, entity: ModelType) -> ModelType: """ Add a new document to the collection. diff --git a/pyfireconsole/queries/query_runner.py b/pyfireconsole/queries/query_runner.py index 68a2807..77d5f51 100644 --- a/pyfireconsole/queries/query_runner.py +++ b/pyfireconsole/queries/query_runner.py @@ -34,6 +34,17 @@ def all(self) -> 'QueryRunner': self.query = AllQuery(self.query or self.collection_key).set_conn(self.conn).exec() return self + def limit(self, limit: int) -> 'QueryRunner': + """ + Apply a limit to the query. + Args: + limit (int): The maximum number of documents to retrieve. + Returns: + QueryRunner: The current QueryRunner instance. + """ + self.query = self.query.limit(limit) if self.query else self.conn.collection(self.collection_key).limit(limit) + return self + def save(self, id: str, data: dict) -> str | None: return SaveQuery(self.collection_key, id, data).set_conn(self.conn).exec() @@ -43,7 +54,10 @@ def create(self, data: dict) -> str | None: def delete(self, id: str) -> None: return DeleteQuery(self.collection_key, id).set_conn(self.conn).exec() - def iter(self) -> Generator[Dict[str, Any], None, None]: + def iter(self, limit: int = 1000) -> Generator[Dict[str, Any], None, None]: docs = self.query.stream() + + docs = self.query.limit(limit).stream() + for doc in docs: yield dict(_doc_to_dict(doc) or {}, id=doc.id) diff --git a/tests/test_usecase.py b/tests/test_usecase.py index fc4bcb9..e849647 100644 --- a/tests/test_usecase.py +++ b/tests/test_usecase.py @@ -205,6 +205,71 @@ def test_first(mock_db): assert user.id in [user1.id, user2.id] +def test_class_all(mock_db): + mock_db = MockFirestore() + mock_db.reset() + + user1 = User.new( + name="John", + email="", + ).save() + user2 = User.new( + name="Mary", + email="", + ).save() + + users = User.all() + + assert len([u for u in users]) == 2 + assert user1.id in [u.id for u in users] + assert user2.id in [u.id for u in users] + + +def test_collection_all(mock_db): + mock_db = MockFirestore() + mock_db.reset() + + user1 = User.new( + name="John", + email="", + ).save() + user2 = User.new( + name="John", + email="", + ).save() + _ = User.new( + name="Mary", + email="", + ).save() + + users = User.where("name", "==", "John").all() + + assert len([u for u in users]) == 2 + assert user1.id in [u.id for u in users] + assert user2.id in [u.id for u in users] + + +def test_to_a(mock_db): + mock_db = MockFirestore() + mock_db.reset() + + user1 = User.new( + name="John", + email="", + ).save() + user2 = User.new( + name="Mary", + email="", + ).save() + + users = User.all().to_a() + + assert len(users) == 2 + assert type(users) is list + assert user1.id in [users[0].id, users[1].id] + assert user2.id in [users[0].id, users[1].id] + + def test_as_json(mock_db): book = Book.new( title="Math", From 2dffa19cffaff88149acf03bea2254354259d290 Mon Sep 17 00:00:00 2001 From: Makoto Tanji Date: Fri, 16 Aug 2024 15:46:30 +0900 Subject: [PATCH 2/2] Version up --- pyproject.toml | 2 +- setup.py | 2 +- 2 files changed, 2 insertions(+), 2 deletions(-) diff --git a/pyproject.toml b/pyproject.toml index 02363e2..4bc313f 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -1,6 +1,6 @@ [tool.poetry] name = "pyfireconsole" -version = "0.1.0" +version = "0.1.1" description = "" authors = ["Makoto Tanji "] readme = "README.md" diff --git a/setup.py b/setup.py index f1ed499..51bd3e9 100644 --- a/setup.py +++ b/setup.py @@ -2,7 +2,7 @@ setup( name='pyfireconsole', - version='0.1.0', + version='0.1.1', author='Makoto Tanji', author_email='tanji.makoto@gmail.com', description='An interactive console for Firestore based on Python ORM',