Use streaming in all assertion comparisons consumers

[Pierre-Sassoulas]

Follow-up to the streaming-comparison chain: #14521 (comparators → generators), #14546 (base comparisons return an iterator), #14587 (set comparison → Iterator[str]), #14588 (PrettyPrinter formats lazily, budget-cappable). Those moved the assertion machinery to generators and a lazy pretty-printer; this PR wires the truncation budget through the remaining consumers so a large diff is never fully built when truncation is going to clip it anyway.

The truncation footer no longer reports the hidden-line count ...Full output truncated (N lines hidden), use '-vv' to show becomes ...Full output truncated, use '-vv' to show. Counting the hidden lines means materialising the whole diff, which is costly. -vv (no truncation) stays exact, and anything within the budget is unchanged.

Benchmark

Small / common case (tiny inputs — the realistic failing assert)

case	9.1.0	this PR	Δ
list [1,2,3] vs [1,2,4]	0.149 ms	0.153 ms	+2.9%
tuple 3 elems	0.145 ms	0.154 ms	+6.1%
set 3 elems	0.155 ms	0.113 ms	-27.0%
dict 3 keys, 1 differ	0.178 ms	0.190 ms	+7.0%
dict extra key	0.110 ms	0.120 ms	+9.1%
str short	0.057 ms	0.061 ms	+7.2%
str 5-line multiline	0.073 ms	0.075 ms	+3.7%
list of 10 dicts	0.269 ms	0.201 ms	-25.3%

Common asserts cost +3% to +9% (a few µs of per-chunk budget bookkeeping); set and list-of-dicts are faster (set-sort fast path + lazy chunked pformat). Absolute cost stays
0.06–0.27 ms.

N = 2000

comparison shape	9.1.0	this PR	Δ	speedup
dict == dict (values differ)	821,887 ms	0.76 ms	-99.9999%	~1,077,000×
str == str (lines differ)	6,953 ms	0.37 ms	-99.995%	~18,800×
list of 3 × 100k-char strings	216.7 ms	0.50 ms	-99.77%	~433×
tuple == tuple	8.54 ms	0.18 ms	-97.8%	~47×
list == list	8.70 ms	0.19 ms	-97.8%	~46×
nested [{N}] == [{}]	17.1 ms	0.49 ms	-97.1%	~35×
dict extra keys ({N}=={})	28.7 ms	0.95 ms	-96.7%	~30×
set == set	9.87 ms	0.42 ms	-95.7%	~23×

N = 20000 (fast shapes; dict==dict values & str==str skipped — O(N²) on main)

comparison shape	9.1.0	this PR	Δ	speedup
tuple == tuple	84.0 ms	0.16 ms	-99.81%	~517×
list == list	84.0 ms	0.16 ms	-99.80%	~511×
nested [{N}] == [{}]	203.7 ms	4.48 ms	-97.8%	~46×
dict extra keys ({N}=={})	352.7 ms	9.32 ms	-97.4%	~38×
set == set	97.9 ms	2.71 ms	-96.9%	~36×

list/tuple stay flat (~0.16 ms) — fully bounded. set (sort O(N log N)), dict extra (nsmallest O(N) scan) and nested still scale with N — far below main, but not constant.
not in (needle not in huge_text) joins the bounded shapes: ~2.8 s → ~3 ms for a 1M-char haystack now that its ndiff inputs are capped.

dict == dict is the extreme: near-identical value lines (1999: 1999, vs 1999: 2000,) trip difflib's intraline _fancy_replace, making the full path ~O(N²). This is
pre-existing — on pytest 8.4.2, pytest -v on assert {i: i for i in range(500)} == {i: i + 1 for i in range(500)} takes 83 s; this PR brings it to ~1 ms.

Benchmark (branch-adaptive; tiers: small / n2000 / n20000)

"""Branch-adaptive benchmark of the assertion-repr pipeline.

Times the full format+truncate path of an assertion explanation for each
comparison shape, at -v (assertion verbosity 1 — the path that both
formats a diff and truncates it), default limits (8 lines / 640 chars),
CI detection off. Adapts to main's API (list + truncate_if_required) vs
the PR (iterator + materialize_with_truncation + pformat_cap).

Tiers: ``small`` (common case), ``n2000``, ``n20000`` (fast shapes only).
Run all, or pass a subset on the command line:

    python bench_assert.py                 # all tiers
    python bench_assert.py small n20000    # skip the slow n2000 tier
"""

import inspect
import os
import sys
import time


os.environ.pop("CI", None)
os.environ.pop("GITHUB_ACTIONS", None)
os.environ.pop("BUILD_NUMBER", None)

import _pytest.assertion.truncate as T  # noqa: E402
import _pytest.assertion.util as U  # noqa: E402
from _pytest.assertion.highlight import dummy_highlighter as H  # noqa: E402


_SIG = inspect.signature(U.assertrepr_compare)
HAS_CAP = "pformat_cap" in _SIG.parameters
HAS_MATERIALIZE = hasattr(T, "materialize_with_truncation")


class FakeConfig:
    def getini(self, name):
        return None  # defaults: 8 lines / 640 chars

    def get_verbosity(self, *a, **k):
        return 1  # -v


class FakeItem:
    config = FakeConfig()


def run(left, right):
    allk = dict(
        op="==",
        left=left,
        right=right,
        verbose=1,
        highlighter=H,
        assertion_text_diff_style=getattr(U, "ASSERTION_TEXT_DIFF_STYLE_NDIFF", "ndiff"),
    )
    if HAS_CAP:
        st, tl, tc = T._get_truncation_parameters(FakeConfig())
        allk["pformat_cap"] = (
            (tl + 3 if tl > 0 else None, tc + 70 if tc > 0 else None)
            if st
            else (None, None)
        )
    kw = {k: v for k, v in allk.items() if k in _SIG.parameters}
    src = U.assertrepr_compare(**kw)
    if HAS_MATERIALIZE:
        return T.materialize_with_truncation(src, FakeConfig())
    return T.truncate_if_required(list(src), FakeItem())


def bench(fn, n):
    return min(
        (lambda: (t := time.perf_counter(), fn(), time.perf_counter() - t)[2])()
        for _ in range(n)
    ) * 1000


def scaled(n):
    return {
        "list == list": (list(range(n)), list(range(1, n + 1))),
        "tuple == tuple": (tuple(range(n)), tuple(range(1, n + 1))),
        "set == set": (set(range(n)), set(range(1, n + 1))),
        "dict diff values": ({i: i for i in range(n)}, {i: i + 1 for i in range(n)}),
        "dict extra ({n}=={})": ({i: i for i in range(n)}, {}),
        "str == str (lines)": ("x\n" * n, "y\n" * n),
        "nested [{n}]==[{}]": ([{i: i for i in range(n)}], [{}]),
    }


# Tiers: name -> {label: (left, right, reps)}
TIERS = {
    "small": {
        "list [1,2,3] vs [1,2,4]": ([1, 2, 3], [1, 2, 4], 20000),
        "tuple 3 elems": ((1, 2, 3), (1, 2, 4), 20000),
        "set 3 elems": ({1, 2, 3}, {1, 2, 4}, 20000),
        "dict 3 keys, 1 differ": (
            {"a": 1, "b": 2, "c": 3},
            {"a": 1, "b": 9, "c": 3},
            20000,
        ),
        "dict extra key": ({"a": 1, "b": 2}, {"a": 1}, 20000),
        "str short": ("hello world", "hello there", 20000),
        "str multiline 5 lines": ("a\nb\nc\nd\ne", "a\nb\nX\nd\ne", 20000),
        "list of 10 dicts": (
            [{"k": i} for i in range(10)],
            [{"k": i + 1} for i in range(10)],
            20000,
        ),
    },
    "n2000": {
        **{k: (left, right, 3) for k, (left, right) in scaled(2000).items()},
        "list of 3 huge strs": (["x" * 100_000] * 3, ["y" * 100_000] * 3, 10),
    },
    # n20000: skip dict==dict (values) and str==str — O(N^2) on main.
    "n20000": {
        k: (left, right, 2)
        for k, (left, right) in scaled(20000).items()
        if k not in ("dict diff values", "str == str (lines)")
    },
}

selected = [a for a in sys.argv[1:] if a in TIERS] or list(TIERS)
api = "stream" if (HAS_CAP and HAS_MATERIALIZE) else "main"
print(f"# API: {api}")
for tier in selected:
    print(f"\n## {tier}")
    print(f"{'case':28}{'ms':>13}  out_lines")
    for label, (left, right, reps) in TIERS[tier].items():
        out = run(left, right)
        ms = bench(lambda left=left, right=right: run(left, right), reps)
        print(f"{label:28}{ms:13.4f}  {len(out)}")

Raw benchmark results:

On main:

~/git/pytest   main  python bench_assert.py

# API: main

## small
case                                   ms  out_lines
list [1,2,3] vs [1,2,4]            0.1476  10
tuple 3 elems                      0.1424  10
set 3 elems                        0.1490  10
dict 3 keys, 1 differ              0.1776  10
dict extra key                     0.1042  10
str short                          0.0545  4
str multiline 5 lines              0.0684  8
list of 10 dicts                   0.2851  10

## n2000
case                                   ms  out_lines
list == list                       9.3782  10
tuple == tuple                     9.1402  10
set == set                         9.2911  10
dict diff values              861306.3471  10
dict extra ({n}=={})              29.0514  10
str == str (lines)              7178.3846  10
nested [{n}]==[{}]                17.6646  5
list of 3 huge strs              221.5416  5

## n20000
case                                   ms  out_lines
list == list                     129.7909  10
tuple == tuple                    94.1027  10
set == set                       104.6896  10
dict extra ({n}=={})             324.3957  10
nested [{n}]==[{}]               179.6054  5

On branch

~/git/pytest   stream-comparisons  python bench_assert.py
# API: stream

## small
case                                   ms  out_lines
list [1,2,3] vs [1,2,4]            0.1598  10
tuple 3 elems                      0.1620  10
set 3 elems                        0.1193  10
dict 3 keys, 1 differ              0.1899  10
dict extra key                     0.1193  10
str short                          0.0608  4
str multiline 5 lines              0.0792  8
list of 10 dicts                   0.2011  10

## n2000
case                                   ms  out_lines
list == list                       0.1897  10
tuple == tuple                     0.1916  10
set == set                         0.4535  10
dict diff values                   0.7497  10
dict extra ({n}=={})               1.0092  10
str == str (lines)                 0.3799  10
nested [{n}]==[{}]                 0.5439  5
list of 3 huge strs                0.5796  5

## n20000
case                                   ms  out_lines
list == list                       0.2136  10
tuple == tuple                     0.3341  10
set == set                         3.5187  10
dict extra ({n}=={})              10.9883  10
nested [{n}]==[{}]                 5.0259  5

[nicoddemus]

Great work!

[Pierre-Sassoulas]

Thank you ! Do you have an opinion about the next step ? (3 options in the PR description)

[bluetech]

See my last comment, I think it might affect the plan, so I'll wait before reviewing the rest.

[bluetech]

Is the annotation needed?

[bluetech]

If we change SetComparisonFunction to return Iterator instead of Iterable, can avoid the iter here. Alternatively, change source to Iterable[str].

[bluetech]

Unfortunately we can't change the return type here, it is stable API (see hookspec.py).

[Pierre-Sassoulas]

Maybe we can truncate in util.assertrepr_compare and still get the performance improvements without modifying the return type of the stable API ?

[bluetech]

I haven't looked at the truncation part, but if you can avoid the API break I'd be happy to take a look.

[Pierre-Sassoulas]

So I moved the truncation to keep the performance without touching the API, but had to hack truncation a little.

First, we have to truncate twice because plugin maintainer expect the truncation to be handled where it was before I suppose.

Second, if we choose to drop the detail in the truncation message ("x lines hidden") then we can also drop the hacky check to see if there's the truncation header in the string for the second truncation. (I would favor that, small price to pay, imo), but right now I don't see an elegant way to do that as we have multiple way to truncate and it's not idempotent.

[nicoddemus]

if we choose to drop the detail in the truncation message ("x lines hidden")

I'm OK with dropping that detail; it is a small price to pay comparing to the benefits.

[Pierre-Sassoulas]

I removed the number of removed line in the message. Benchmark on pathological case with very big set show big improvement, but I benchmarked on pylint test suite for a realistic perf change:

round1-upstream: 118.29s
round1-HEAD : 115.39s
round2-upstream: 114.87s
round2-HEAD : 106.43s
round3-upstream: 116.51s
round3-HEAD : 118.69s
round4-upstream: 114.14s
round4-HEAD : 116.28s

Look like it's within noise. pylint does not compare big data structure (the bigger checks are probably the functional test output ~=80 lines).