From a8e3fd4759330403a808207e993861626c5bbb93 Mon Sep 17 00:00:00 2001
From: Manish Yadav <ymanish@ariqt.com>
Date: Tue, 1 Sep 2026 13:26:05 +0530
Subject: [PATCH] fix: get_leaf_nodes KeyError on nodes-less leaf nodes

clean_node() deletes the `nodes` key entirely from childless nodes
during tree cleanup instead of setting it to an empty list, so
get_leaf_nodes()'s direct `structure['nodes']` access raised KeyError
on any tree produced by the normal pipeline (#330).

Co-Authored-By: Claude Sonnet 5 <noreply@anthropic.com>
---
 pageindex/utils.py  |  2 +-
 tests/test_utils.py | 48 +++++++++++++++++++++++++++++++++++++++++++++
 2 files changed, 49 insertions(+), 1 deletion(-)
 create mode 100644 tests/test_utils.py

diff --git a/pageindex/utils.py b/pageindex/utils.py
index 83e9d0e79..d3f7ceb91 100644
--- a/pageindex/utils.py
+++ b/pageindex/utils.py
@@ -300,7 +300,7 @@ def structure_to_list(structure):
     
 def get_leaf_nodes(structure):
     if isinstance(structure, dict):
-        if not structure['nodes']:
+        if not structure.get('nodes'):
             structure_node = copy.deepcopy(structure)
             structure_node.pop('nodes', None)
             return [structure_node]
diff --git a/tests/test_utils.py b/tests/test_utils.py
new file mode 100644
index 000000000..a683430ce
--- /dev/null
+++ b/tests/test_utils.py
@@ -0,0 +1,48 @@
+from pageindex.utils import get_leaf_nodes
+
+
+def test_get_leaf_nodes_handles_missing_nodes_key():
+    """Tree cleanup (see `clean_node` in utils.py) deletes the `nodes` key
+    entirely from childless nodes rather than setting it to `[]`, so
+    `get_leaf_nodes` must not assume the key is always present (#330)."""
+    tree = {
+        "title": "Root",
+        "nodes": [
+            {"title": "Leaf A"},  # no 'nodes' key at all
+            {
+                "title": "Branch",
+                "nodes": [
+                    {"title": "Leaf B"},
+                ],
+            },
+        ],
+    }
+
+    leaves = get_leaf_nodes(tree)
+
+    assert {leaf["title"] for leaf in leaves} == {"Leaf A", "Leaf B"}
+    assert all("nodes" not in leaf for leaf in leaves)
+
+
+def test_get_leaf_nodes_still_handles_empty_nodes_list():
+    tree = {
+        "title": "Root",
+        "nodes": [
+            {"title": "Leaf A", "nodes": []},
+        ],
+    }
+
+    leaves = get_leaf_nodes(tree)
+
+    assert [leaf["title"] for leaf in leaves] == ["Leaf A"]
+
+
+def test_get_leaf_nodes_handles_list_of_trees():
+    trees = [
+        {"title": "Leaf A"},
+        {"title": "Root", "nodes": [{"title": "Leaf B"}]},
+    ]
+
+    leaves = get_leaf_nodes(trees)
+
+    assert {leaf["title"] for leaf in leaves} == {"Leaf A", "Leaf B"}