tree

Author: njzjz

deepmd/dpmodel/utils/serialization.py

@@ -337,12 +337,17 @@ def from_list(cls, data_list: list[Any]) -> "Node":
         return cls("ListNode", children, data, variables)
 
     def __str__(self) -> str:
+        elbow = "└──"
+        pipe = "│  "
+        tee = "├──"
+        blank = "   "
+        linebreak = "\n"
         buff = []
         buff.append(f"{self.name} (size={format_big_number(self.size)})")
         children_buff = []
-        for kk, vv in self.children.items():
+        for ii, (kk, vv) in enumerate(self.children.items()):
             # add indentation
-            child_repr = str(vv).replace("\n", "\n  ")
+            child_repr = str(vv)
             if len(children_buff) > 0:
                 # check if it is the same as the last one
                 last_repr = children_buff[-1][1]
@@ -360,5 +365,13 @@ def format_list_keys(kk: str) -> str:
                     return f"[{keys[0]}...{keys[-1]}]"
             return kk
 
-        buff.extend(f"  {format_list_keys(kk)} -> {vv}" for kk, vv in children_buff)
+        buff.extend(
+            f"{tee if ii < len(children_buff) - 1 else elbow}{format_list_keys(kk)} -> {
+                vv.replace(
+                    linebreak,
+                    linebreak + (pipe if ii < len(children_buff) - 1 else blank),
+                )
+            }"
+            for ii, (kk, vv) in enumerate(children_buff)
+        )
         return "\n".join(buff)

source/tests/common/dpmodel/test_network.py

@@ -312,8 +312,8 @@ def test_save_load_model_yaml(self) -> None:
     def test_node_display(self):
         disp_expected = textwrap.dedent("""\
             some_class some_type (size=18)
-              layers -> ListNode (size=18)
-                0, 1 -> Node (size=9)""")
+            └──layers -> ListNode (size=18)
+               └──0, 1 -> Node (size=9)""")
         disp = str(Node.deserialize(self.model_dict))
         self.assertEqual(disp, disp_expected)