support named tuples

Author: samuelcolvin

devtools/prettier.py

@@ -40,7 +40,8 @@ def __init__(self,
             (dict, self._format_dict),
             (str, self._format_str),
             (bytes, self._format_bytes),
-            ((tuple, list, set, frozenset), self._format_list_like),
+            (tuple, self._format_tuples),
+            ((list, set, frozenset), self._format_list_like),
             (collections.Generator, self._format_generators),
         ]
 
@@ -95,6 +96,22 @@ def _format_list_like(self, value: Union[list, tuple, set], value_repr: str, ind
             self._stream.write(',\n')
         self._stream.write(indent_current * self._c + close_)
 
+    def _format_tuples(self, value: tuple, value_repr: str, indent_current: int, indent_new: int):
+        fields = getattr(value, '_fields', None)
+        if fields:
+            # named tuple
+            self._stream.write(value.__class__.__name__ + '(\n')
+            for field, v in zip(fields, value):
+                self._stream.write(indent_new * self._c)
+                self._stream.write(field)
+                self._stream.write('=')
+                self._format(v, indent_new, False)
+                self._stream.write(',\n')
+            self._stream.write(indent_current * self._c + ')')
+        else:
+            # normal tuples are just like other similar iterables
+            return self._format_list_like(value, value_repr, indent_current, indent_new)
+
     def _format_str(self, value: str, value_repr: str, indent_current: int, indent_new: int):
         if self._repr_strings:
             self._stream.write(value_repr)

tests/test_prettier.py

@@ -1,5 +1,6 @@
 import string
 from collections import OrderedDict
+from typing import Any, NamedTuple
 
 import numpy
 
@@ -82,6 +83,19 @@ def test_generator():
         ')')
 
 
+def test_named_tuple():
+    class TestNT(NamedTuple):
+        foo: str
+        bar: str
+        spam: int
+    v = pformat(TestNT('x', 'y', 1))
+    assert v == ("TestNT(\n"
+                 "    foo='x',\n"
+                 "    bar='y',\n"
+                 "    spam=1,\n"
+                 ")")
+
+
 def test_generator_no_yield():
     pformat_ = PrettyFormat(yield_from_generators=False)
     v = pformat_((i for i in range(3)))
@@ -158,3 +172,47 @@ def test_frozenset():
     1,
     2,
 })"""
+
+
+def test_deep_objects():
+    class TestNT(NamedTuple):
+        foo: str
+        bar: str
+        spam: Any
+    v = pformat((
+        {
+            'a': TestNT('x', 'y', OrderedDict([(1, 2), (3, 4), (5, 6)])),
+            'b': frozenset(range(3)),
+            'c': [1, 2, {1: 2, 3: 4}]
+        },
+        {1, 2, 3}
+    ))
+    print(v)
+    assert v == """\
+(
+    {
+        'a': TestNT(
+            foo='x',
+            bar='y',
+            spam=OrderedDict([
+                (1, 2),
+                (3, 4),
+                (5, 6),
+            ]),
+        ),
+        'b': frozenset({
+            0,
+            1,
+            2,
+        }),
+        'c': [
+            1,
+            2,
+            {
+                1: 2,
+                3: 4,
+            },
+        ],
+    },
+    {1, 2, 3},
+)"""