File tree Expand file tree Collapse file tree 2 files changed +33
-0
lines changed Expand file tree Collapse file tree 2 files changed +33
-0
lines changed Original file line number Diff line number Diff line change 4
4
"""
5
5
import textwrap
6
6
from typing import (
7
+ Generator ,
7
8
NamedTuple ,
8
9
Sequence ,
9
10
Tuple ,
@@ -248,6 +249,19 @@ def pretty(
248
249
)
249
250
return text
250
251
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
+
251
265
# NOTE:
252
266
# The values of the properties defined below directly map to properties
253
267
# of the underlying `Token`s. A root node does not translate to a `Token`
Original file line number Diff line number Diff line change @@ -72,3 +72,22 @@ def test_pretty(file_regression):
72
72
)
73
73
node = SyntaxTreeNode (tokens )
74
74
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
You can’t perform that action at this time.
0 commit comments