[perf] Make PrettyPrinter format lazily so output can be budget-capped

[Pierre-Sassoulas]

Refactor required prior to #14523.

_format and the per-type helpers now yield their output as a stream of string chunks instead of writing to a file-like object, and pformat joins them. On top of that, pformat_lines pulls from the formatter only until a budget is reached:

pformat_lines(obj, max_lines=None, max_chars=None)

It stops on the first chunk that reaches either budget, so a huge collection costs O(budget) rather than O(N). Either dimension may be None (unbounded); with both None the whole object is formatted.

Benchmark (PrettyPrinter alone, width 80)::

list(range(500_000)):
    pformat().splitlines()        ~805 ms
    pformat_lines(max_lines=11)   ~0.027 ms      (~30000x)

[8 small ints] (common small diff):
    pformat().splitlines()        ~0.0133 ms
    pformat_lines(max_lines=11)   ~0.0163 ms (+3µs)

["x"*100_000] * 3 (flat, few huge elements):
    pformat_lines(max_chars=640)  stops after ~100_000 chars
                                  (one element) instead of 300_000

[bluetech]

Thanks, seems like a nice optimization to me, and I'd say makes the code nicer as well.

[bluetech]

Probably a "public" method shouldn't reference a private method in its docstring. I would describe the behavior rather than the implementation here.

[bluetech]

Would remove the detail of what the caller does, or replace "truncates" with "can truncate" or "should truncate".

[bluetech]

I suggest enforcing max_lines and max_chars are kw-only. Otherwise they can be easily confused.

[Pierre-Sassoulas]

Yes, definitely.

[bluetech]

This is an unrelated optimization, or is it somehow related to the iterator change?

If it's an optimization, a comment would be helpful, something like: "Try direct sort first, faster than the fallback.".

[Pierre-Sassoulas]

Yes sorry I didn't realize this was in this commit. It's indeed an optimization (rather big one imho). Because it uses C sort directly it's a lot faster, and also I suppose that we rarely have heterogenous data structure. (Then it's a little longer of course). Here's a micro-benchmark:

Details

"""Interleaved benchmark of the set-sort fast path in ``_pprint_set``.

Compares ``sorted(s)`` against ``sorted(s, key=_safe_key)`` (the prior
behaviour), and the heterogeneous fallback ``try sorted / except retry``
against always using ``_safe_key``.

Interleaving: A and B are timed back-to-back every iteration, so slow
drift (CPU frequency scaling, thermal, GC) hits both equally and cancels
in the per-iteration difference. We report the median per-call time of
each and the median of the *paired* differences (B - A), which is the
robust estimate of the real gap.
"""

import statistics
import time

from _pytest._io.pprint import _safe_key


def bench_pair(a, b, inner, outer):
    """Interleave-time ``a`` and ``b``; return per-call medians + paired diff."""
    ta_samples, tb_samples, diffs = [], [], []
    for _ in range(outer):
        t = time.perf_counter()
        for _ in range(inner):
            a()
        ta = (time.perf_counter() - t) / inner
        t = time.perf_counter()
        for _ in range(inner):
            b()
        tb = (time.perf_counter() - t) / inner
        ta_samples.append(ta)
        tb_samples.append(tb)
        diffs.append(tb - ta)
    return (
        statistics.median(ta_samples) * 1000,  # ms
        statistics.median(tb_samples) * 1000,  # ms
        statistics.median(diffs) * 1e6,  # us, paired
    )


print("plain sorted (A) vs sorted(key=_safe_key) (B), homogeneous:")
for label, data, inner, outer in [
    ("int set 1k", set(range(1000)), 100, 300),
    ("int set 100k", set(range(100_000)), 3, 80),
    ("str set 100k", {f"item-{i}" for i in range(100_000)}, 1, 80),
]:
    a = lambda d=data: sorted(d)
    b = lambda d=data: sorted(d, key=_safe_key)
    ma, mb, diff = bench_pair(a, b, inner, outer)
    print(
        f"  {label:13} A={ma:9.4f} ms  B={mb:9.4f} ms  B/A={mb / ma:5.1f}x"
        f"  paired B-A={diff:+10.1f} us"
    )

print("\nheterogeneous (unorderable mix), fallback path:")
het = {1, "a", 2, "b", 3.5, None} | {f"x{i}" for i in range(500)} | set(range(500))
try:
    sorted(het)
    raise SystemExit("expected TypeError - set is not heterogeneous")
except TypeError:
    pass


def safe_sort(d=het):
    return sorted(d, key=_safe_key)


def failed_sort(d=het):
    # the *only* extra work the try/except adds: a plain sort that raises
    try:
        sorted(d)
    except TypeError:
        pass


# Measure the overhead directly: the failed sort is the whole cost of the
# fallback's try/except. Comparing new(=failed+safe) vs old(=safe) is
# useless here — the ~us signal is far below the jitter of the ~ms safe
# sort, so the paired diff is noise (can even come out negative).
overhead_ms, safe_ms, _ = bench_pair(failed_sort, safe_sort, 20, 300)
print(
    f"  {'hetero 1k':13} _safe_key sort={safe_ms:8.4f} ms"
    f"  try/except overhead (failed sort)={overhead_ms * 1000:6.1f} us"
    f"  = {overhead_ms / safe_ms * 100:.2f}% of the sort"
)

Result on my machine:

plain sorted (A) vs sorted(key=_safe_key) (B), homogeneous:
  int set 1k    A=   0.0153 ms  B=   0.4661 ms  B/A= 30.4x  paired B-A=    +451.4 us
  int set 100k  A=   1.8300 ms  B=  63.6966 ms  B/A= 34.8x  paired B-A=  +61787.4 us
  str set 100k  A=  54.3415 ms  B= 356.0647 ms  B/A=  6.6x  paired B-A= +302171.5 us

heterogeneous (unorderable mix), fallback path:
  hetero 1k     _safe_key sort=  4.1314 ms  try/except overhead (failed sort)=  12.2 us  = 0.30% of the sort

[bluetech]

Might as well assert it as well.