Skip to content

Commit 60c8ff8

Browse files
authored
feat(sdk): node navigation helpers get_node, get_node_parent, get_node_path, get_node_map (#542)
They take the tree, not a doc_id: fetch it once (get_document_structure) and navigate locally, instead of one request per step. get_node / get_node_parent / get_node_path share one depth-first walk, O(n) per call. They return the tree's own nodes rather than copies (unlike get_nodes / get_leaf_nodes), so node["nodes"] keeps working. get_node_map is create_node_mapping plus the input check; create_node_mapping stays as the 0.2.8 surface. Wrong input raises TypeError naming get_document_structure(doc_id) / get_tree(doc_id)['result'] instead of reading as "not found": the whole get_tree() response, a None tree, or a non-string id. create_node_mapping (so get_node_map too) walks past a node whose nodes is None instead of raising TypeError, as get_node already did.
1 parent f279431 commit 60c8ff8

2 files changed

Lines changed: 69 additions & 2 deletions

File tree

‎pageindex/utils.py‎

Lines changed: 35 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1257,7 +1257,7 @@ def create_node_mapping(tree, include_page_ranges=False, max_page=None):
12571257
"end_index"} (end = next node's page_index, or max_page for the last node)."""
12581258
def get_all_nodes(tree):
12591259
if isinstance(tree, dict):
1260-
return [tree] + [node for child in tree.get('nodes', []) for node in get_all_nodes(child)]
1260+
return [tree] + [node for child in tree.get('nodes') or [] for node in get_all_nodes(child)]
12611261
elif isinstance(tree, list):
12621262
return [node for item in tree for node in get_all_nodes(item)]
12631263
return []
@@ -1276,6 +1276,40 @@ def get_all_nodes(tree):
12761276
}
12771277
return mapping
12781278

1279+
def _require_node_tree(tree):
1280+
if not (isinstance(tree, list) or (isinstance(tree, dict) and 'node_id' in tree)):
1281+
raise TypeError("tree must be a node list such as get_document_structure(doc_id) "
1282+
"or get_tree(doc_id)['result'], not the whole get_tree response; "
1283+
f"got {type(tree).__name__}")
1284+
1285+
def get_node_path(tree, node_id):
1286+
"""[top-level ancestor, ..., node] for node_id; [] if absent."""
1287+
_require_node_tree(tree)
1288+
if not isinstance(node_id, str):
1289+
raise TypeError(f"node_id must be a str like '0007', got {node_id!r}")
1290+
for node in [tree] if isinstance(tree, dict) else tree:
1291+
if node.get('node_id') == node_id:
1292+
return [node]
1293+
path = get_node_path(node.get('nodes') or [], node_id)
1294+
if path:
1295+
return [node] + path
1296+
return []
1297+
1298+
def get_node(tree, node_id):
1299+
"""The node with node_id, or None."""
1300+
path = get_node_path(tree, node_id)
1301+
return path[-1] if path else None
1302+
1303+
def get_node_parent(tree, node_id):
1304+
"""The parent of node_id; None for a top-level or absent node."""
1305+
path = get_node_path(tree, node_id)
1306+
return path[-2] if len(path) > 1 else None
1307+
1308+
def get_node_map(tree):
1309+
"""{node_id: node} for every node in tree."""
1310+
_require_node_tree(tree)
1311+
return create_node_mapping(tree)
1312+
12791313
def print_tree(tree, exclude_fields=None, indent=0):
12801314
"""Outline view; passing exclude_fields gives the 0.2.8 pprint view."""
12811315
if exclude_fields is not None:

‎tests/test_package_surface.py‎

Lines changed: 34 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -3,7 +3,11 @@
33
import subprocess
44
import sys
55

6-
from pageindex.utils import create_node_mapping, print_tree, remove_fields
6+
import pytest
7+
8+
from pageindex.utils import (create_node_mapping, get_node, get_node_map,
9+
get_node_parent, get_node_path, print_tree,
10+
remove_fields)
711

812
TREE = [
913
{"title": "Root", "node_id": "0000", "page_index": 1,
@@ -47,6 +51,35 @@ def test_print_tree_exclude_fields(capsys):
4751
assert "[0000] Root" in capsys.readouterr().out
4852

4953

54+
# ── tree navigation ──
55+
56+
def test_node_navigation():
57+
root, child, tail = TREE[0], TREE[0]["nodes"][0], TREE[1]
58+
assert get_node(TREE, "0001") is child
59+
assert get_node(TREE, "9999") is None
60+
assert get_node_parent(TREE, "0001") is root
61+
assert get_node_parent(TREE, "0002") is None
62+
assert get_node_path(TREE, "0001") == [root, child]
63+
assert get_node_path(TREE, "0002") == [tail]
64+
assert get_node_path(TREE, "9999") == []
65+
assert get_node(root, "0001") is child
66+
assert get_node_map(TREE) == {"0000": root, "0001": child, "0002": tail}
67+
leaf = {"node_id": "0000", "nodes": None}
68+
assert get_node_map([leaf]) == {"0000": leaf}
69+
70+
71+
def test_node_navigation_rejects_wrong_input():
72+
envelope = {"doc_id": "d", "status": "completed", "result": TREE}
73+
with pytest.raises(TypeError, match="got dict"):
74+
get_node(envelope, "0001")
75+
with pytest.raises(TypeError, match="got dict"):
76+
get_node_map(envelope)
77+
with pytest.raises(TypeError, match="got NoneType"):
78+
get_node_map(None)
79+
with pytest.raises(TypeError, match="node_id"):
80+
get_node(TREE, 1)
81+
82+
5083
# ── import cost: the SDK must not pay for the indexing stack ──
5184

5285
def test_import_pageindex_is_lazy():

0 commit comments

Comments
 (0)