Fix recursive_stringify

src/robofish/io/file.py

@@ -627,7 +627,7 @@ class File(h5py.File):
         """
         return robofish.io.validate(self, strict_validate)
 
-    def to_string(self, output_format: str = "shape") -> str:
+    def to_string(self, output_format: str = "shape", max_width: int = 120) -> str:
         """The file is formatted to a human readable format.
         Args:
             output_format: ['shape', 'full'] show the shape, or the full content of datasets
@@ -636,46 +636,111 @@ class File(h5py.File):
         """
 
         def recursive_stringify(
-            obj: h5py.Group, output_format: str, level: int = 0
+            obj: h5py.Group,
+            output_format: str,
+            parent_indices: List[int] = [],
+            parent_siblings: List[int] = [],
         ) -> str:
             """This function crawls recursively into hdf5 groups.
             Datasets and attributes are directly attached, for groups, the function is recursively called again.
             Args:
                 obj: a h5py group
                 output_format: ['shape', 'full'] show the shape, or the full content of datasets
-                level: the current indentation level
             Returns:
                 A string representation of the group
             """
+
+            def lines(dataset_attribute: bool = False) -> str:
+                """Get box-drawing characters for the graph lines."""
+                line = ""
+                for pi, ps in zip(parent_indices, parent_siblings):
+                    if pi < ps - 1:
+                        line += "│ "
+                    else:
+                        line += "  "
+                if dataset_attribute:
+                    line += "  "
+                line += "─ "
+                junction_index = 2 * len(parent_indices) + dataset_attribute * 2 - 1
+                last = "└"
+                other = "├"
+                if dataset_attribute:
+                    j = (
+                        last
+                        if list(value.attrs.keys()).index(d_key) == len(value.attrs) - 1
+                        else other
+                    )
+                else:
+                    j = last if index == num_children - 1 else other
+                line = line[: junction_index + 1] + j + line[junction_index + 1 :]
+                if isinstance(value, h5py.Group) or (
+                    isinstance(value, h5py.Dataset)
+                    and not dataset_attribute
+                    and value.attrs
+                ):
+                    line = line[:-1] + "┬─"
+                else:
+                    line = line[:-1] + "──"
+
+                return line + " "
+
             s = ""
-            level_str = "|---" * level
-            for key, value in obj.attrs.items():
-                s += "%s %s:\t%s\n" % (level_str, key, value)
+            max_key_len = 0
+            num_children = 0
+            if obj.attrs:
+                max_key_len = max(len(key) for key in obj.attrs)
+                num_children += len(obj.attrs)
+            if hasattr(obj, "items"):
+                max_key_len = max([len(key) for key in obj] + [max_key_len])
+                num_children += len(obj)
+            index = 0
+            if obj.attrs:
+                for key, value in obj.attrs.items():
+                    value = str(value).replace("\n", " ").strip()
+                    if len(value) > max_width - max_key_len - len(lines()):
+                        value = (
+                            value[: max_width - max_key_len - len(lines()) - 3] + "..."
+                        )
+                    s += f"{lines()}{key: <{max_key_len}}  {value}\n"
+                    index += 1
             if hasattr(obj, "items"):
                 for key, value in obj.items():
                     if isinstance(value, h5py.Dataset):
                         if output_format == "shape":
-                            s += "%s %s:\t Shape %s\n" % (level_str, key, value.shape)
+                            s += (
+                                f"{lines()}"
+                                f"{key: <{max_key_len}}  Shape {value.shape}\n"
+                            )
                         else:
-                            s += "%s %s:\n%s\n" % (
-                                level_str,
-                                key,
-                                np.array2string(
-                                    value,
-                                    precision=2,
-                                    separator=" ",
-                                    suppress_small=True,
-                                ),
+                            s += f"{lines()}{key}:\n"
+                            s += np.array2string(
+                                value,
+                                precision=2,
+                                separator=" ",
+                                suppress_small=True,
                             )
+                            s += "\n"
 
+                        if value.attrs:
+                            d_max_key_len = max(len(dk) for dk in value.attrs)
                         for d_key, d_value in value.attrs.items():
-                            s += "%s %s:\t%s\n" % (level_str + "|---", d_key, d_value)
+                            d_value = str(d_value).replace("\n", " ").strip()
+                            if len(d_value) > max_width - d_max_key_len - len(
+                                lines(True)
+                            ):
+                                d_value = d_value[
+                                    : max_width - d_max_key_len - len(lines(True))
+                                ]
+                                d_value = d_value[:-3] + "..."
+                            s += f"{lines(True)}{d_key: <{d_max_key_len}}  {d_value}\n"
                     if isinstance(value, h5py.Group):
-                        s += "%s| %s\n%s" % (
-                            level_str,
-                            key,
-                            recursive_stringify(value, output_format, level + 1),
+                        s += f"{lines()}{key}\n" + recursive_stringify(
+                            obj=value,
+                            output_format=output_format,
+                            parent_indices=parent_indices + [index],
+                            parent_siblings=parent_siblings + [num_children],
                         )
+                    index += 1
             return s
 
         return recursive_stringify(self, output_format)