PEP 813: keep Expand.ALWAYS through hugging, and support compact in every mode

Lib/_pprint_protocol.py

@@ -503,7 +503,7 @@ def _layout_document(document, level, options):
                 hugged = _hugged_index(items)
             parts = [_item_document(item, level + 1, options) for item in items]
             if hugged is not None:
-                return _hugged_document(header, parts)
+                return _hugged_document(header, parts, level, options)
             positional = not any(
                 isinstance(item, (_EntryDocument, _KeywordDocument)) for item in items
             )
@@ -609,31 +609,58 @@ def _hugged_index(items):
     return found
 
 
-def _hugged_document(header, parts):
+def _hugged_document(header, parts, level, options):
     """Keep every item on the opening and closing lines.
 
     Only the one container among the items breaks across lines, so its
     brackets hug the object's own, as in ``defaultdict(<class 'list'>, {``.
     """
     close = header.single_close if len(parts) == 1 else header.close
     separator = _Text(header.sep + _SPACE)
-    return _Group(
-        _Concat((_Text(header.open), _join(separator, parts), _Text(close)))
-    )
+    doc = _Concat((_Text(header.open), _join(separator, parts), _Text(close)))
+    if options.expand is Expand.ALWAYS and level > 0:
+        return doc
+    return _Group(doc)
+
+
+def _compact_body(separator, parts, column, options):
+    """Join items, packing those that fit on one line, as PrettyPrinter does.
+
+    An item whose one-line form fits in the width left after *column* shares
+    lines with its packable neighbors.  Any other item gets lines of its own,
+    so a container that breaks never starts or ends partway along a line.
+    """
+    available = options.width - column
+    packable = []
+    for part in parts:
+        lines = list(_render_lines(_Group(part), _sys.maxsize))
+        packable.append(len(lines) == 1 and len(lines[0]) < available)
+    docs = [parts[0]]
+    for index in range(1, len(parts)):
+        if packable[index - 1] and packable[index]:
+            line = _Group(_Concat((_Line(), parts[index])))
+        else:
+            line = _Concat((_Line(), parts[index]))
+        docs.append(_Concat((separator, line)))
+    return _Concat(tuple(docs))
 
 
 def _container_document(header, parts, positional, level, options):
     close = header.single_close if len(parts) == 1 else header.close
     separator = _Text(header.sep)
+    compact = options.compact and positional
 
     if options.expand is Expand.NEVER:
-        if options.compact and positional:
-            body = _join(_Concat((separator, _Group(_Line()))), parts)
-            return _Concat((_Text(header.open), _Align(body), _Text(close)))
-        body = _join(_Concat((separator, _Line())), parts)
+        if compact:
+            body = _compact_body(separator, parts, level + 1, options)
+        else:
+            body = _join(_Concat((separator, _Line())), parts)
         return _Group(_Concat((_Text(header.open), _Align(body), _Text(close))))
 
-    body = _join(_Concat((separator, _Line())), parts)
+    if compact:
+        body = _compact_body(separator, parts, options.indent * (level + 1), options)
+    else:
+        body = _join(_Concat((separator, _Line())), parts)
     doc = _Concat(
         (
             _Text(header.open),

Lib/test/test__pprint.py

@@ -490,6 +490,92 @@ def test_hug_ignored_when_never(self):
             "defaultdict(<class 'list'>,\n"
         ))
 
+    def test_always_expands_inside_hug(self):
+        obj = [list(range(4)), collections.deque([[1, 2], [3]])]
+        self.assertEqual(
+            pformat(obj, width=25, expand=Expand.ALWAYS),
+            "[\n"
+            "    [\n"
+            "        0,\n"
+            "        1,\n"
+            "        2,\n"
+            "        3,\n"
+            "    ],\n"
+            "    deque([\n"
+            "        [\n"
+            "            1,\n"
+            "            2,\n"
+            "        ],\n"
+            "        [\n"
+            "            3,\n"
+            "        ],\n"
+            "    ]),\n"
+            "]",
+        )
+
+
+class CompactTests(unittest.TestCase):
+    def test_fit_packs_items(self):
+        self.assertEqual(
+            pformat([(1, 2)] * 8, width=30, compact=True),
+            "[\n"
+            "    (1, 2), (1, 2), (1, 2),\n"
+            "    (1, 2), (1, 2), (1, 2),\n"
+            "    (1, 2), (1, 2),\n"
+            "]",
+        )
+
+    def test_fit_broken_item_has_own_lines(self):
+        self.assertEqual(
+            pformat([1, 2, list(range(12)), 3, 4], width=20, compact=True),
+            "[\n"
+            "    1, 2,\n"
+            "    [\n"
+            "        0, 1, 2, 3,\n"
+            "        4, 5, 6, 7,\n"
+            "        8, 9, 10,\n"
+            "        11,\n"
+            "    ],\n"
+            "    3, 4,\n"
+            "]",
+        )
+
+    def test_always_does_not_pack_containers(self):
+        self.assertEqual(
+            pformat(
+                [list(range(12)), [1]],
+                width=20, compact=True, expand=Expand.ALWAYS,
+            ),
+            "[\n"
+            "    [\n"
+            "        0, 1, 2, 3,\n"
+            "        4, 5, 6, 7,\n"
+            "        8, 9, 10,\n"
+            "        11,\n"
+            "    ],\n"
+            "    [\n"
+            "        1,\n"
+            "    ],\n"
+            "]",
+        )
+
+    def test_never_matches_pretty_printer(self):
+        obj = [1, 2, list(range(12)), 3, 4]
+        self.assertEqual(
+            pformat(obj, width=20, compact=True, expand=Expand.NEVER, indent=1),
+            pprint.PrettyPrinter(width=20, compact=True).pformat(obj),
+        )
+
+    def test_key_value_items_are_not_packed(self):
+        self.assertEqual(
+            pformat({"a": 1, "b": 2, "c": 3}, width=10, compact=True),
+            "{\n"
+            "    'a': 1,\n"
+            "    'b': 2,\n"
+            "    'c': 3,\n"
+            "}",
+        )
+
 
 class DescriptionTests(unittest.TestCase):
     def test_ordered_dict(self):