diff --git a/frappe/model/document.py b/frappe/model/document.py index d0ad8106fa..2a969fb910 100644 --- a/frappe/model/document.py +++ b/frappe/model/document.py @@ -167,7 +167,7 @@ def get_docs( as_iterator: bool = False, for_update: bool = False, distinct: bool = False, -) -> list["Document"] | Generator[list["Document"]]: +) -> list["Document"] | Generator["Document"]: """Fetch fully instantiated Document objects from the database. Returns a list of Documents by default. Pass `as_iterator=True` to get @@ -175,7 +175,7 @@ def get_docs( :param doctype: DocType of the records to fetch. :param filters: Dict or list of filters to apply. - :param chunk_size: Number of records to yield per chunk if using `as_iterator`. Default 1000. + :param chunk_size: Number of records to fetch in each chunk if using `as_iterator`. :param limit: Maximum total number of records to fetch. :param limit_start: Start results at record #. Default 0. :param order_by: Order By string, e.g. `creation desc`. @@ -243,7 +243,7 @@ def _get_docs_generator( lock_rows, for_update, distinct, -) -> Generator[list["Document"]]: +) -> Generator["Document"]: fetched_count = 0 current_offset = limit_start @@ -270,7 +270,7 @@ def _get_docs_generator( break built_docs = _build_document_objects(controller, chunk_data, for_update) - yield built_docs + yield from built_docs fetched_count += len(chunk_data) current_offset += len(chunk_data) diff --git a/frappe/tests/test_document.py b/frappe/tests/test_document.py index 478d84cf68..94e2f0e643 100644 --- a/frappe/tests/test_document.py +++ b/frappe/tests/test_document.py @@ -903,28 +903,11 @@ class TestGetDocs(IntegrationTestCase): docs_desc = frappe.get_docs(self.parent_dt, order_by="creation desc") self.assertEqual(docs_asc[0].name, docs_desc[-1].name) - def test_generator_yields_chunks(self): - chunks = list(frappe.get_docs(self.parent_dt, as_iterator=True, chunk_size=2)) - # 5 records with chunk_size=2 should give 3 chunks (2, 2, 1) - self.assertEqual(len(chunks), 3) - self.assertEqual(len(chunks[0]), 2) - self.assertEqual(len(chunks[1]), 2) - self.assertEqual(len(chunks[2]), 1) - - def test_generator_with_limit(self): - chunks = list(frappe.get_docs(self.parent_dt, as_iterator=True, chunk_size=2, limit=3)) - total = sum(len(c) for c in chunks) - self.assertEqual(total, 3) - def test_generator_parity(self): eager = frappe.get_docs(self.parent_dt, order_by="creation asc") - gen_docs = [ - doc - for chunk in frappe.get_docs( - self.parent_dt, as_iterator=True, chunk_size=2, order_by="creation asc" - ) - for doc in chunk - ] + gen_docs = list( + frappe.get_docs(self.parent_dt, as_iterator=True, chunk_size=2, order_by="creation asc") + ) self.assertEqual([d.name for d in eager], [d.name for d in gen_docs]) def test_for_update_sets_flag(self):