Skip to content

Commit 7a1df9c

Browse files
authored
✨ NEW: Add SyntaxTreeNode.walk (#148)
1 parent 2fbc419 commit 7a1df9c

File tree

2 files changed

+33
-0
lines changed

2 files changed

+33
-0
lines changed

markdown_it/tree.py

Lines changed: 14 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -4,6 +4,7 @@
44
"""
55
import textwrap
66
from typing import (
7+
Generator,
78
NamedTuple,
89
Sequence,
910
Tuple,
@@ -248,6 +249,19 @@ def pretty(
248249
)
249250
return text
250251

252+
def walk(
253+
self: _NodeType, *, include_self: bool = True
254+
) -> Generator[_NodeType, None, None]:
255+
"""Recursively yield all descendant nodes in the tree starting at self.
256+
257+
The order mimics the order of the underlying linear token
258+
stream (i.e. depth first).
259+
"""
260+
if include_self:
261+
yield self
262+
for child in self.children:
263+
yield from child.walk(include_self=True)
264+
251265
# NOTE:
252266
# The values of the properties defined below directly map to properties
253267
# of the underlying `Token`s. A root node does not translate to a `Token`

tests/test_tree.py

Lines changed: 19 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -72,3 +72,22 @@ def test_pretty(file_regression):
7272
)
7373
node = SyntaxTreeNode(tokens)
7474
file_regression.check(node.pretty(indent=2, show_text=True), extension=".xml")
75+
76+
77+
def test_walk():
78+
tokens = MarkdownIt().parse(EXAMPLE_MARKDOWN)
79+
tree = SyntaxTreeNode(tokens)
80+
expected_node_types = (
81+
"root",
82+
"heading",
83+
"inline",
84+
"text",
85+
"paragraph",
86+
"inline",
87+
"text",
88+
"strong",
89+
"text",
90+
"text",
91+
)
92+
for node, expected_type in zip(tree.walk(), expected_node_types):
93+
assert node.type == expected_type

0 commit comments

Comments
 (0)