Flatten GroupNodes for ASCII tree.

tests/test_model.py

@@ -41,15 +41,12 @@ tree = """+--rw (test:choiA)?
 |  +--rw testb:leafT?
 |  +--rw listA* [leafE leafF]
 |     +--rw contD
-|     |  +---x acA
-|     |  |  +--ro input
-|     |  |  +--ro output
-|     |  |     +--ro leafL
 |     |  +--rw contE!
 |     |  |  +--rw leafJ?
 |     |  |  +--rw leafP?
 |     |  |  +--rw leafU?
 |     |  +--rw leafG?
+|     |  +--ro leafL
 |     +--rw leafE
 |     +--rw leafF
 +--rw test:contT
@@ -67,13 +64,9 @@ tree = """+--rw (test:choiA)?
 |  +--rw uint32?
 |  +--rw uint64?
 |  +--rw uint8?
-+---n testb:noA
-|  +--rw testb:leafO?
-+---x testb:rpcA
-   +--ro testb:input
-   |  +--ro testb:leafK?
-   +--ro testb:output
-      +--ro testb:llistC*
++--ro testb:leafK?
++--rw testb:leafO?
++--ro testb:llistC*
 """
 
 @pytest.fixture
@@ -195,6 +188,7 @@ def test_schema(data_model):
     assert data_model.get_data_node("/testb:noA/leafO") is None
 
 def test_tree(data_model):
+    print(data_model.ascii_tree())
     assert data_model.ascii_tree() == tree
 
 def test_types(data_model):

yangson/schema.py

@@ -81,6 +81,9 @@ class SchemaNode:
         """
         return [r.data_path() for r in self._state_roots()]
 
+    def _flatten(self) -> List["SchemaNode"]:
+        return [self]
+
     def _handle_substatements(self, stmt: Statement, mid: ModuleId) -> None:
         """Dispatch actions for substatements of `stmt`."""
         for s in stmt.substatements:
@@ -453,12 +456,16 @@ class InternalNode(SchemaNode):
 
     def _ascii_tree(self, indent: str) -> str:
         """Return the receiver's subtree as ASCII art."""
+        if not self.children: return ""
+        cs = []
+        for c in self.children:
+            cs.extend(c._flatten())
+        cs.sort(key=lambda x: x.qual_name)
         res = ""
-        if not self.children: return res
-        cn = sorted(self.children, key=lambda x: x.qual_name)
-        for c in cn[:-1]:
+        for c in cs[:-1]:
             res += indent + c._tree_line() + c._ascii_tree(indent + "|  ")
-        return res + indent + cn[-1]._tree_line() + cn[-1]._ascii_tree(indent + "   ")
+        return (res + indent + cs[-1]._tree_line() +
+                cs[-1]._ascii_tree(indent + "   "))
 
 class GroupNode(InternalNode):
     """Anonymous group of schema nodes."""
@@ -473,6 +480,12 @@ class GroupNode(InternalNode):
     def _pattern_entry(self) -> SchemaPattern:
         return super()._schema_pattern()
 
+    def _flatten(self) -> List["SchemaNode"]:
+        res = []
+        for c in self.children:
+            res.extend(c._flatten())
+        return res
+
 class DataNode(SchemaNode):
     """Abstract superclass for data nodes."""