fix: use pages dict for page break generation to support failed pages

Only generate page breaks for pages present in DoclingDocument.pages dict.
This enables proper page break markers for failed pages (added by docling)
while maintaining compatibility with filter() method (which removes pages).

Changes:
- Add page_numbers parameter to _yield_page_breaks() function
- Extract page_numbers from doc.pages.keys() in _iterate_items()
- Update test data to include failed pages in pages dict
- Update test expectations for new behavior

Signed-off-by: jhchoi1182 <jhchoi1182@gmail.com>

Author: jhchoi1182

docling_core/transforms/serializer/common.py

@@ -84,23 +84,34 @@ def _yield_page_breaks(
     next_page: int,
     lvl: int,
     start_index: int,
+    page_numbers: Optional[set[int]] = None,
 ) -> Iterable[tuple[_PageBreakNode, int, int]]:
     """Yield page break nodes for each page in range (prev_page, next_page].
 
     Generates one PageBreakNode per page transition. For example, if prev_page=1
-    and next_page=4, yields 3 page breaks for pages 2, 3, and 4.
+    and next_page=4, and page_numbers contains pages 1, 2, 3, 4, yields 3 page
+    breaks for pages 2, 3, and 4.
+
+    If page_numbers is provided, only generates page breaks for pages that exist
+    in page_numbers. This ensures filtered documents (via filter()) don't generate
+    spurious page breaks for excluded pages.
 
     Args:
         prev_page: The last seen page number (1-based physical index).
         next_page: The current page number (1-based physical index).
         lvl: The nesting level for the yielded nodes.
         start_index: The starting index for page break node IDs.
+        page_numbers: Optional set of valid page numbers. If provided, only pages
+            in this set will generate page breaks.
 
     Yields:
         Tuples of (PageBreakNode, level, next_index) for each page transition.
     """
     idx = start_index
     for page in range(prev_page + 1, next_page + 1):
+        # Skip pages that are not in the document's pages dict
+        if page_numbers is not None and page not in page_numbers:
+            continue
         yield (
             _PageBreakNode(
                 self_ref=f"#/pb/{idx}",
@@ -124,6 +135,9 @@ def _iterate_items(
     my_visited: set[str] = visited if visited is not None else set()
     prev_page_nr: Optional[int] = None
     page_break_i = 0
+    # Get the set of valid page numbers from the document's pages dict
+    # This ensures filtered documents don't generate spurious page breaks
+    page_numbers: set[int] = set(doc.pages.keys())
     for item, lvl in doc.iterate_items(
         root=node,
         with_groups=True,
@@ -146,7 +160,7 @@ def _iterate_items(
                         page_no = it.prov[0].page_no
                         if prev_page_nr is not None and page_no > prev_page_nr:
                             for pb_node, pb_lvl, page_break_i in _yield_page_breaks(
-                                prev_page_nr, page_no, lvl, page_break_i
+                                prev_page_nr, page_no, lvl, page_break_i, page_numbers
                             ):
                                 yield pb_node, pb_lvl
                         # update previous page number to avoid duplicate page breaks
@@ -157,7 +171,7 @@ def _iterate_items(
                 if prev_page_nr is None or page_no > prev_page_nr:
                     if prev_page_nr is not None:  # close previous range
                         for pb_node, pb_lvl, page_break_i in _yield_page_breaks(
-                            prev_page_nr, page_no, lvl, page_break_i
+                            prev_page_nr, page_no, lvl, page_break_i, page_numbers
                         ):
                             yield pb_node, pb_lvl
                     prev_page_nr = page_no

test/data/doc/2408.09869v3_enriched_p2_p3_p5.gt.html

@@ -72,10 +72,7 @@
             "b": 672.1780283203125,
             "coord_origin": "BOTTOMLEFT"
           },
-          "charspan": [
-            0,
-            30
-          ]
+          "charspan": [0, 30]
         }
       ],
       "orig": "78 ICT\u00a0\uae30\uc220\ubcc0\ud654\uc5d0\u00a0\ub530\ub978\u00a0\ubbf8\ub798\u00a0\ubcf4\uc548\uae30\uc220\u00a0\uc804\ub9dd\u00a0\ubcf4\uace0\uc11c",
@@ -99,10 +96,7 @@
             "b": 615.8270283203125,
             "coord_origin": "BOTTOMLEFT"
           },
-          "charspan": [
-            0,
-            3
-          ]
+          "charspan": [0, 3]
         }
       ],
       "orig": "\ubcf4\uc548\ubd84",
@@ -127,10 +121,7 @@
             "b": 672.1780283203125,
             "coord_origin": "BOTTOMLEFT"
           },
-          "charspan": [
-            0,
-            30
-          ]
+          "charspan": [0, 30]
         }
       ],
       "orig": "84 ICT\u00a0\uae30\uc220\ubcc0\ud654\uc5d0\u00a0\ub530\ub978\u00a0\ubbf8\ub798\u00a0\ubcf4\uc548\uae30\uc220\u00a0\uc804\ub9dd\u00a0\ubcf4\uace0\uc11c",
@@ -154,10 +145,7 @@
             "b": 359.95302832031246,
             "coord_origin": "BOTTOMLEFT"
           },
-          "charspan": [
-            0,
-            164
-          ]
+          "charspan": [0, 164]
         }
       ],
       "orig": "- \u00a0 AI\u00a0\uae30\uc220\uc744\u00a0\ud65c\uc6a9\ud574\u00a0\uc624\ud0d0\uc728\uc744\u00a0\uc904\uc774\uace0,\u00a0\ube60\ub978\u00a0\ubd84\uc11d\u00a0\ud0d0\uc9c0\ub97c\u00a0\ub118\uc5b4\u00a0\uc790\ub3d9\uc801\uc778\u00a0\uaca9\ub9ac\uae4c\uc9c0\u00a0\uc81c\uacf5\ud558\ub294\u00a0\ubc29\ud5a5\uc73c\ub85c\u00a0\uc9c4\ud654 - \u00a0 \uac01\uc885\u00a0\uc0ac\uc774\ubc84\uc704\ud611\uc5d0\u00a0\ub300\ud55c\u00a0\ud558\ub098\uc758\u00a0\uae30\ubc95\u00a0\ub610\ub294\u00a0\uc54c\uace0\ub9ac\uc998\uc774\u00a0\uc544\ub2cc,\u00a0\uce68\ud574\u00a0\ub300\uc751\uc758\u00a0\uac01\u00a0\uc694\uc18c\ubcc4\ub85c\u00a0\ucd5c\uc801\ud654\ub41c\u00a0\uc54c\uace0\ub9ac\uc998\uacfc \uae30\ubc95,\u00a0\uc774\uc804\uc5d0\u00a0\uc124\uba85\ud55c\u00a0\ud559\uc2b5\u00a0\ubc29\ubc95\uc744\u00a0\uc9c0\uc6d0\ud558\uace0\u00a0\uc774\ub97c\u00a0\ud1b5\ud569\ud560\u00a0\uc218\u00a0\uc788\ub294\u00a0\uae30\uc220\u00a0\ud544\uc694",
@@ -185,10 +173,7 @@
             "b": 620.6187973022461,
             "coord_origin": "BOTTOMLEFT"
           },
-          "charspan": [
-            0,
-            0
-          ]
+          "charspan": [0, 0]
         }
       ],
       "captions": [],
@@ -208,6 +193,13 @@
       },
       "page_no": 1
     },
+    "2": {
+      "size": {
+        "width": 515.906005859375,
+        "height": 728.5040283203125
+      },
+      "page_no": 2
+    },
     "3": {
       "size": {
         "width": 515.906005859375,
@@ -216,4 +208,4 @@
       "page_no": 3
     }
   }
-}
+}

test/data/doc/skipped_1page.json

@@ -72,10 +72,7 @@
             "b": 672.1780283203125,
             "coord_origin": "BOTTOMLEFT"
           },
-          "charspan": [
-            0,
-            30
-          ]
+          "charspan": [0, 30]
         }
       ],
       "orig": "78 ICT\u00a0\uae30\uc220\ubcc0\ud654\uc5d0\u00a0\ub530\ub978\u00a0\ubbf8\ub798\u00a0\ubcf4\uc548\uae30\uc220\u00a0\uc804\ub9dd\u00a0\ubcf4\uace0\uc11c",
@@ -99,10 +96,7 @@
             "b": 615.8270283203125,
             "coord_origin": "BOTTOMLEFT"
           },
-          "charspan": [
-            0,
-            3
-          ]
+          "charspan": [0, 3]
         }
       ],
       "orig": "\ubcf4\uc548\ubd84",
@@ -127,10 +121,7 @@
             "b": 672.1780283203125,
             "coord_origin": "BOTTOMLEFT"
           },
-          "charspan": [
-            0,
-            30
-          ]
+          "charspan": [0, 30]
         }
       ],
       "orig": "84 ICT\u00a0\uae30\uc220\ubcc0\ud654\uc5d0\u00a0\ub530\ub978\u00a0\ubbf8\ub798\u00a0\ubcf4\uc548\uae30\uc220\u00a0\uc804\ub9dd\u00a0\ubcf4\uace0\uc11c",
@@ -154,10 +145,7 @@
             "b": 359.95302832031246,
             "coord_origin": "BOTTOMLEFT"
           },
-          "charspan": [
-            0,
-            164
-          ]
+          "charspan": [0, 164]
         }
       ],
       "orig": "- \u00a0 AI\u00a0\uae30\uc220\uc744\u00a0\ud65c\uc6a9\ud574\u00a0\uc624\ud0d0\uc728\uc744\u00a0\uc904\uc774\uace0,\u00a0\ube60\ub978\u00a0\ubd84\uc11d\u00a0\ud0d0\uc9c0\ub97c\u00a0\ub118\uc5b4\u00a0\uc790\ub3d9\uc801\uc778\u00a0\uaca9\ub9ac\uae4c\uc9c0\u00a0\uc81c\uacf5\ud558\ub294\u00a0\ubc29\ud5a5\uc73c\ub85c\u00a0\uc9c4\ud654 - \u00a0 \uac01\uc885\u00a0\uc0ac\uc774\ubc84\uc704\ud611\uc5d0\u00a0\ub300\ud55c\u00a0\ud558\ub098\uc758\u00a0\uae30\ubc95\u00a0\ub610\ub294\u00a0\uc54c\uace0\ub9ac\uc998\uc774\u00a0\uc544\ub2cc,\u00a0\uce68\ud574\u00a0\ub300\uc751\uc758\u00a0\uac01\u00a0\uc694\uc18c\ubcc4\ub85c\u00a0\ucd5c\uc801\ud654\ub41c\u00a0\uc54c\uace0\ub9ac\uc998\uacfc \uae30\ubc95,\u00a0\uc774\uc804\uc5d0\u00a0\uc124\uba85\ud55c\u00a0\ud559\uc2b5\u00a0\ubc29\ubc95\uc744\u00a0\uc9c0\uc6d0\ud558\uace0\u00a0\uc774\ub97c\u00a0\ud1b5\ud569\ud560\u00a0\uc218\u00a0\uc788\ub294\u00a0\uae30\uc220\u00a0\ud544\uc694",
@@ -185,10 +173,7 @@
             "b": 620.6191101074219,
             "coord_origin": "BOTTOMLEFT"
           },
-          "charspan": [
-            0,
-            0
-          ]
+          "charspan": [0, 0]
         }
       ],
       "captions": [],
@@ -208,6 +193,20 @@
       },
       "page_no": 1
     },
+    "2": {
+      "size": {
+        "width": 515.906005859375,
+        "height": 728.5040283203125
+      },
+      "page_no": 2
+    },
+    "3": {
+      "size": {
+        "width": 515.906005859375,
+        "height": 728.5040283203125
+      },
+      "page_no": 3
+    },
     "4": {
       "size": {
         "width": 515.906005859375,
@@ -216,4 +215,4 @@
       "page_no": 4
     }
   }
-}
+}

test/data/doc/skipped_2pages.json

@@ -35,26 +35,28 @@ def test_normal_4pages_has_all_pages():
     ], f"Expected pages [1, 2, 3, 4], got {page_numbers}"
 
 
-def test_skipped_2pages_has_only_two_pages():
-    """Test that skipped_2pages.json has only 2 pages (pages 2 and 3 failed to parse)."""
+def test_skipped_2pages_has_all_pages_including_failed():
+    """Test that skipped_2pages.json has all 4 pages (pages 2 and 3 failed to parse but are still present)."""
     src = Path("./test/data/doc/skipped_2pages.json")
     doc = DoclingDocument.load_from_json(src)
 
     page_numbers = list(doc.pages.keys())
 
-    assert len(page_numbers) == 2, f"Expected 2 pages in skipped_2pages.json, got {len(page_numbers)}"
-    assert page_numbers == [1, 4], f"Expected pages [1, 4], got {page_numbers}"
+    # After fix: all pages including failed ones should be present in pages dict
+    assert len(page_numbers) == 4, f"Expected 4 pages in skipped_2pages.json, got {len(page_numbers)}"
+    assert page_numbers == [1, 2, 3, 4], f"Expected pages [1, 2, 3, 4], got {page_numbers}"
 
 
-def test_skipped_1page_has_two_pages():
-    """Test that skipped_1page.json has 2 pages (page 2 failed to parse)."""
+def test_skipped_1page_has_all_pages_including_failed():
+    """Test that skipped_1page.json has all 3 pages (page 2 failed to parse but is still present)."""
     src = Path("./test/data/doc/skipped_1page.json")
     doc = DoclingDocument.load_from_json(src)
 
     page_numbers = list(doc.pages.keys())
 
-    assert len(page_numbers) == 2, f"Expected 2 pages in skipped_1page.json, got {len(page_numbers)}"
-    assert page_numbers == [1, 3], f"Expected pages [1, 3], got {page_numbers}"
+    # After fix: all pages including failed ones should be present in pages dict
+    assert len(page_numbers) == 3, f"Expected 3 pages in skipped_1page.json, got {len(page_numbers)}"
+    assert page_numbers == [1, 2, 3], f"Expected pages [1, 2, 3], got {page_numbers}"
 
 
 # =============================================================================