Pretty-printing types (#1430)

* Pass 'indent' and 'compact' down Type._str.

* Types now have a 'show' command (that indents).

src/awkward/_v2/types/arraytype.py

@@ -1,5 +1,7 @@
 # BSD 3-Clause License; see https://github.com/scikit-hep/awkward-1.0/blob/main/LICENSE
 
+import sys
+
 import awkward as ak
 import awkward._v2.types.type
 from awkward._v2.forms.form import _parameters_equal
@@ -35,7 +37,13 @@ def length(self):
         return self._length
 
     def __str__(self):
-        return str(self._length) + " * " + str(self._content)
+        return "".join(self._str("", True))
+
+    def show(self, stream=sys.stdout):
+        stream.write("".join(self._str("", False) + ["\n"]))
+
+    def _str(self, indent, compact):
+        return [str(self._length) + " * "] + self._content._str(indent, compact)
 
     def __repr__(self):
         args = [repr(self._content), repr(self._length)]

src/awkward/_v2/types/listtype.py

@@ -39,24 +39,26 @@ def __init__(self, content, parameters=None, typestr=None):
     def content(self):
         return self._content
 
-    def __str__(self):
+    def _str(self, indent, compact):
         if self._typestr is not None:
-            out = self._typestr
+            out = [self._typestr]
 
         elif self.parameter("__array__") == "string":
-            return "string"
+            return ["string"]
 
         elif self.parameter("__array__") == "bytestring":
-            return "bytes"
+            return ["bytes"]
 
         else:
             params = self._str_parameters()
             if params is None:
-                out = f"var * {str(self._content)}"
+                out = ["var * "] + self._content._str(indent, compact)
             else:
-                out = f"[var * {str(self._content)}, {params}]"
+                out = (
+                    ["[var * "] + self._content._str(indent, compact) + [f", {params}]"]
+                )
 
-        return self._str_categorical_begin() + out + self._str_categorical_end()
+        return [self._str_categorical_begin()] + out + [self._str_categorical_end()]
 
     def __repr__(self):
         args = [repr(self._content)] + self._repr_args()

src/awkward/_v2/types/numpytype.py

@@ -122,15 +122,15 @@ def primitive(self):
 
     _str_parameters_exclude = ("__categorical__", "__unit__")
 
-    def __str__(self):
+    def _str(self, indent, compact):
         if self._typestr is not None:
-            out = self._typestr
+            out = [self._typestr]
 
         elif self.parameter("__array__") == "char":
-            out = "char"
+            out = ["char"]
 
         elif self.parameter("__array__") == "byte":
-            out = "byte"
+            out = ["byte"]
 
         else:
             if self.parameter("__unit__") is not None:
@@ -143,17 +143,17 @@ def __str__(self):
             params = self._str_parameters()
 
             if units is None and params is None:
-                out = self._primitive
+                out = [self._primitive]
             else:
                 if units is not None and params is not None:
                     units = units + ", "
                 elif units is None:
                     units = ""
                 elif params is None:
                     params = ""
-                out = self._primitive + "[" + units + params + "]"
+                out = [self._primitive, "[", units, params, "]"]
 
-        return self._str_categorical_begin() + out + self._str_categorical_end()
+        return [self._str_categorical_begin()] + out + [self._str_categorical_end()]
 
     def __repr__(self):
         args = [repr(self._primitive)] + self._repr_args()

src/awkward/_v2/types/optiontype.py

@@ -42,21 +42,23 @@ def __init__(self, content, parameters=None, typestr=None):
     def content(self):
         return self._content
 
-    def __str__(self):
+    def _str(self, indent, compact):
         if self._typestr is not None:
-            out = self._typestr
+            out = [self._typestr]
 
         else:
             params = self._str_parameters()
             if params is None:
                 if not isinstance(self._content, (RegularType, ListType)):
-                    out = f"?{str(self._content)}"
+                    out = ["?"] + self._content._str(indent, compact)
                 else:
-                    out = f"option[{str(self._content)}]"
+                    out = ["option["] + self._content._str(indent, compact) + ["]"]
             else:
-                out = f"option[{str(self._content)}, {params}]"
+                out = (
+                    ["option["] + self._content._str(indent, compact) + [f", {params}]"]
+                )
 
-        return self._str_categorical_begin() + out + self._str_categorical_end()
+        return [self._str_categorical_begin()] + out + [self._str_categorical_end()]
 
     def __repr__(self):
         args = [repr(self._content)] + self._repr_args()

src/awkward/_v2/types/recordtype.py

@@ -2,9 +2,8 @@
 
 from collections.abc import Iterable
 
-import json
-
 import awkward as ak
+import awkward._v2._prettyprint
 from awkward._v2.types.type import Type
 from awkward._v2.forms.form import _parameters_equal
 
@@ -73,53 +72,80 @@ def is_tuple(self):
 
     _str_parameters_exclude = ("__categorical__", "__record__")
 
-    def __str__(self):
+    def _str(self, indent, compact):
         if self._typestr is not None:
-            out = self._typestr
+            out = [self._typestr]
 
         else:
-            children = [str(x) for x in self._contents]
+            if compact:
+                pre, post = "", ""
+            else:
+                pre, post = "\n" + indent + "    ", "\n" + indent
+
+            children = []
+            for i, x in enumerate(self._contents):
+                if i + 1 < len(self._contents):
+                    if compact:
+                        y = x._str(indent, compact) + [", "]
+                    else:
+                        y = x._str(indent + "    ", compact) + [",\n", indent, "    "]
+                else:
+                    if compact:
+                        y = x._str(indent, compact)
+                    else:
+                        y = x._str(indent + "    ", compact)
+                children.append(y)
+
             params = self._str_parameters()
             name = self.parameter("__record__")
 
+            if not self.is_tuple:
+                pairs = []
+                for k, v in zip(self._fields, children):
+                    if ak._v2._prettyprint.is_identifier.match(k) is None:
+                        key_str = repr(k)
+                        if key_str.startswith("u"):
+                            key_str = key_str[1:]
+                    else:
+                        key_str = k
+                    pairs.append([key_str, ": "] + v)
+                flat_pairs = [y for x in pairs for y in x]
+
             if params is None:
                 if self.is_tuple:
+                    flat_children = [y for x in children for y in x]
                     if name is None:
-                        out = "(" + ", ".join(children) + ")"
+                        out = ["(", pre] + flat_children + [post, ")"]
                     else:
-                        out = name + "[" + ", ".join(children) + "]"
+                        out = [name, "[", pre] + flat_children + [post, "]"]
                 else:
-                    pairs = []
-                    for k, v in zip(self._fields, children):
-                        if ak._v2._prettyprint.is_identifier.match(k) is None:
-                            key_str = repr(k)
-                            if key_str.startswith("u"):
-                                key_str = key_str[1:]
-                        else:
-                            key_str = k
-                        pairs.append(key_str + ": " + v)
                     if name is None:
-                        out = "{" + ", ".join(pairs) + "}"
+                        out = ["{", pre] + flat_pairs + [post, "}"]
                     else:
-                        out = name + "[" + ", ".join(pairs) + "]"
+                        out = [name, "[", pre] + flat_pairs + [post, "]"]
 
             else:
                 if self.is_tuple:
+                    flat_children = [y for x in children for y in x]
                     if name is None:
-                        out = "tuple[[{}], {}]".format(", ".join(children), params)
+                        out = (
+                            ["tuple[[", pre]
+                            + flat_children
+                            + [post, "], ", params, "]"]
+                        )
                     else:
-                        out = "{}[{}, {}]".format(name, ", ".join(children), params)
+                        out = (
+                            [name, "[", pre] + flat_children + [", ", post, params, "]"]
+                        )
                 else:
                     if name is None:
-                        fields = [json.dumps(x) for x in self._fields]
-                        out = "struct[[{}], [{}], {}]".format(
-                            ", ".join(fields), ", ".join(children), params
+                        out = (
+                            ["struct[{", pre] + flat_pairs + [post, "}, ", params, "]"]
                         )
                     else:
-                        pairs = [k + ": " + v for k, v in zip(self._fields, children)]
-                        out = "{}[{}, {}]".format(name, ", ".join(pairs), params)
+                        out = [name, "[", pre] + flat_pairs + [", ", post, params, "]"]
 
-        return self._str_categorical_begin() + out + self._str_categorical_end()
+        return [self._str_categorical_begin()] + out + [self._str_categorical_end()]
 
     def __repr__(self):
         args = [repr(self._contents), repr(self._fields)] + self._repr_args()

src/awkward/_v2/types/regulartype.py

@@ -52,24 +52,28 @@ def content(self):
     def size(self):
         return self._size
 
-    def __str__(self):
+    def _str(self, indent, compact):
         if self._typestr is not None:
-            out = self._typestr
+            out = [self._typestr]
 
         elif self.parameter("__array__") == "string":
-            return f"string[{self._size}]"
+            return [f"string[{self._size}]"]
 
         elif self.parameter("__array__") == "bytestring":
-            return f"bytes[{self._size}]"
+            return [f"bytes[{self._size}]"]
 
         else:
             params = self._str_parameters()
             if params is None:
-                out = f"{self._size} * {str(self._content)}"
+                out = [str(self._size), " * "] + self._content._str(indent, compact)
             else:
-                out = f"[{self._size} * {str(self._content)}, {params}]"
+                out = (
+                    ["[", str(self._size), " * "]
+                    + self._content._str(indent, compact)
+                    + [", ", params, "]"]
+                )
 
-        return self._str_categorical_begin() + out + self._str_categorical_end()
+        return [self._str_categorical_begin()] + out + [self._str_categorical_end()]
 
     def __repr__(self):
         args = [repr(self._content), repr(self._size)] + self._repr_args()

src/awkward/_v2/types/type.py

@@ -1,6 +1,7 @@
 # BSD 3-Clause License; see https://github.com/scikit-hep/awkward-1.0/blob/main/LICENSE
 
 import json
+import sys
 
 import awkward as ak
 
@@ -22,6 +23,12 @@ def parameter(self, key):
     def typestr(self):
         return self._typestr
 
+    def __str__(self):
+        return "".join(self._str("", True))
+
+    def show(self, stream=sys.stdout):
+        stream.write("".join(self._str("", False) + ["\n"]))
+
     _str_parameters_exclude = ("__categorical__",)
 
     def _str_categorical_begin(self):

src/awkward/_v2/types/uniontype.py

@@ -52,20 +52,39 @@ def __init__(self, contents, parameters=None, typestr=None):
     def contents(self):
         return self._contents
 
-    def __str__(self):
+    def _str(self, indent, compact):
         if self._typestr is not None:
-            out = self._typestr
+            out = [self._typestr]
 
         else:
-            children = [str(x) for x in self._contents]
+            if compact:
+                pre, post = "", ""
+            else:
+                pre, post = "\n" + indent + "    ", "\n" + indent
+
+            children = []
+            for i, x in enumerate(self._contents):
+                if i + 1 < len(self._contents):
+                    if compact:
+                        y = x._str(indent, compact) + [", "]
+                    else:
+                        y = x._str(indent + "    ", compact) + [",\n", indent, "    "]
+                else:
+                    if compact:
+                        y = x._str(indent, compact)
+                    else:
+                        y = x._str(indent + "    ", compact)
+                children.append(y)
+
+            flat_children = [y for x in children for y in x]
             params = self._str_parameters()
 
             if params is None:
-                out = "union[{}]".format(", ".join(children))
+                out = ["union[", pre] + flat_children + [post, "]"]
             else:
-                out = "union[{}, {}]".format(", ".join(children), params)
+                out = ["union[", pre] + flat_children + [", ", post, params, "]"]
 
-        return self._str_categorical_begin() + out + self._str_categorical_end()
+        return [self._str_categorical_begin()] + out + [self._str_categorical_end()]
 
     def __repr__(self):
         args = [repr(self._contents)] + self._repr_args()

src/awkward/_v2/types/unknowntype.py

@@ -26,18 +26,18 @@ def __init__(self, parameters=None, typestr=None):
         self._parameters = parameters
         self._typestr = typestr
 
-    def __str__(self):
+    def _str(self, indent, compact):
         if self._typestr is not None:
-            out = self._typestr
+            out = [self._typestr]
 
         else:
             params = self._str_parameters()
             if params is None:
-                out = "unknown"
+                out = ["unknown"]
             else:
-                out = "unknown[" + params + "]"
+                out = ["unknown[", params, "]"]
 
-        return self._str_categorical_begin() + out + self._str_categorical_end()
+        return [self._str_categorical_begin()] + out + [self._str_categorical_end()]
 
     def __repr__(self):
         args = self._repr_args()

tests/v2/test_0914-types-and-forms.py

@@ -660,7 +660,7 @@ def test_RecordType():
                 {"x": 123},
             )
         )
-        == 'struct[["x", "y"], [unknown, bool], parameters={"x": 123}]'
+        == 'struct[{x: unknown, y: bool}, parameters={"x": 123}]'
     )
     assert (
         str(
@@ -876,7 +876,7 @@ def test_RecordType():
                 {"x": 123, "__categorical__": True},
             )
         )
-        == 'categorical[type=struct[["x", "y"], [unknown, bool], parameters={"x": 123}]]'
+        == 'categorical[type=struct[{x: unknown, y: bool}, parameters={"x": 123}]]'
     )
     assert (
         str(