Use streaming in all assertion comparisons consumers

changelog/14523.improvement.rst

@@ -0,0 +1,13 @@
+Assertion explanations are now built lazily and the truncator stops
+the comparison helpers as soon as it has enough output, so comparing
+two large collections no longer builds the full diff in order to
+discard it. A focused micro-benchmark the worst case scenario
+(``set(range(500_000)) == set(range(1, 500_001))``) drops from ~2,200 ms
+to ~43 ms; but realistic test suite with mostly small diffs should be
+unchanged.
+
+The truncation footer no longer reports the hidden-line count
+(``...Full output truncated (N lines hidden), ...`` becomes
+``...Full output truncated, ...``); diff lines now carry a redundant
+``\x1b[0m`` reset prefix (invisible to terminals) so we can handle
+line one by one.

doc/en/example/reportingdemo.rst

@@ -148,7 +148,7 @@ Here is a nice run of several failures and how ``pytest`` presents things:
     E           1
     E           1...
     E
-    E         ...Full output truncated (7 lines hidden), use '-vv' to show
+    E         ...Full output truncated, use '-vv' to show
 
     failure_demo.py:62: AssertionError
     _________________ TestSpecialisedExplanations.test_eq_list _________________

doc/en/how-to/output.rst

@@ -172,7 +172,7 @@ Now we can increase pytest's verbosity:
     E               'banana',
     E               'apple',...
     E
-    E         ...Full output truncated (7 lines hidden), use '-vv' to show
+    E         ...Full output truncated, use '-vv' to show
 
     test_verbosity_example.py:8: AssertionError
     ____________________________ test_numbers_fail _____________________________
@@ -190,7 +190,7 @@ Now we can increase pytest's verbosity:
     E         {'10': 10, '20': 20, '30': 30, '40': 40}
     E         ...
     E
-    E         ...Full output truncated (16 lines hidden), use '-vv' to show
+    E         ...Full output truncated, use '-vv' to show
 
     test_verbosity_example.py:14: AssertionError
     ___________________________ test_long_text_fail ____________________________

src/_pytest/_io/pprint.py

@@ -56,6 +56,51 @@ def _safe_tuple(t):
     return _safe_key(t[0]), _safe_key(t[1])
 
 
+class _LineBudgetExceeded(Exception):
+    """Internal: signals that ``_LineBudgetStream`` has reached its cap.
+
+    Raised from inside ``stream.write`` so the formatter's recursion
+    unwinds at the next write; the caller catches it and uses whatever
+    output accumulated so far.
+    """
+
+
+class _LineBudgetStream:
+    r"""Stream that collects ``write`` calls into lines and bails out
+    once ``max_lines`` have been produced.
+
+    Lets ``PrettyPrinter._format`` stop early on huge collections: the
+    formatter writes one ``\n``-terminated chunk per element, so the
+    budget check fires on element boundaries.
+    """
+
+    __slots__ = ("_lines", "_max", "_pending")
+
+    def __init__(self, max_lines: int) -> None:
+        self._lines: list[str] = []
+        self._pending: list[str] = []
+        self._max = max_lines
+
+    def write(self, s: str) -> None:
+        if "\n" not in s:
+            if s:
+                self._pending.append(s)
+            return
+        parts = s.split("\n")
+        self._pending.append(parts[0])
+        self._lines.append("".join(self._pending))
+        self._lines.extend(parts[1:-1])
+        self._pending = [parts[-1]] if parts[-1] else []
+        if len(self._lines) >= self._max:
+            raise _LineBudgetExceeded
+
+    def finish(self) -> list[str]:
+        if self._pending:
+            self._lines.append("".join(self._pending))
+            self._pending = []
+        return self._lines
+
+
 class PrettyPrinter:
     def __init__(
         self,
@@ -91,6 +136,48 @@ def pformat(self, object: Any) -> str:
         self._format(object, sio, 0, 0, set(), 0)
         return sio.getvalue()
 
+    def pformat_lines(self, object: Any, max_lines: int | None = None) -> list[str]:
+        """Pretty-print ``object`` and return its lines.
+
+        With ``max_lines=None`` this is just ``self.pformat(object).
+        splitlines()`` — the fast C path through ``StringIO`` and
+        ``str.splitlines``.
+
+        With ``max_lines`` set, the formatter is wired to a Python-level
+        stream that raises once that many lines have been produced; the
+        caller gets back whatever was emitted before the abort. For
+        huge collections this turns an O(N) pformat into O(``max_lines``)
+        — useful when a downstream truncator is going to drop everything
+        past a small budget anyway.
+
+        Sized containers small enough that pformat will obviously fit
+        in the budget take the fast path too: the pure-Python budget
+        stream is ~1.3x slower than ``StringIO`` + ``splitlines`` on
+        tiny inputs (per-write overhead dominates), and paying that on
+        common-case small assertion diffs is a net loss.
+        """
+        if max_lines is None:
+            return self.pformat(object).splitlines()
+        # Sufficient condition for "the budget will never fire": one
+        # pformat line per element is the lower bound for any
+        # container, so ``len(object) <= max_lines`` guarantees the
+        # full pformat fits in the budget.
+        try:
+            size = len(object)
+        except TypeError:
+            size = -1
+        if 0 <= size <= max_lines:
+            return self.pformat(object).splitlines()
+        stream = _LineBudgetStream(max_lines)
+        try:
+            # ``_format``'s ``IO[str]`` annotation is overly tight — it
+            # only ever calls ``stream.write(str)``, which is all this
+            # budget-aware stream implements.
+            self._format(object, stream, 0, 0, set(), 0)  # type: ignore[arg-type]
+        except _LineBudgetExceeded:
+            pass
+        return stream.finish()
+
     def _format(
         self,
         object: Any,
@@ -236,7 +323,12 @@ def _pprint_set(
         else:
             stream.write(typ.__name__ + "({")
             endchar = "})"
-        object = sorted(object, key=_safe_key)
+        try:
+            object = sorted(object)
+        except TypeError:
+            # Heterogeneous element types — fall back to a key that
+            # tolerates unorderable pairs by string-comparing their types.
+            object = sorted(object, key=_safe_key)
         self._format_items(object, stream, indent, allowance, context, level)
         stream.write(endchar)
 

src/_pytest/assertion/__init__.py

@@ -181,13 +181,21 @@ def callbinrepr(op, left: object, right: object) -> str | None:
             config=item.config, op=op, left=left, right=right
         )
         for new_expl in hook_result:
+            # Plugin-supplied lists are truncated here; the built-in impl
+            # already truncates as it streams, so re-applying truncation
+            # to its output is a near no-op (the body fits the budget,
+            # only the footer line is re-emitted with the same wording).
+            # ``materialize_with_truncation`` can return ``[]`` when the
+            # input was a truthy-but-empty iterable, so re-check after
+            # materialising.
             if new_expl:
-                new_expl = truncate.truncate_if_required(new_expl, item)
-                new_expl = [line.replace("\n", "\\n") for line in new_expl]
-                res = "\n~".join(new_expl)
-                if item.config.getvalue("assertmode") == "rewrite":
-                    res = res.replace("%", "%%")
-                return res
+                new_expl = truncate.materialize_with_truncation(new_expl, item.config)
+                if new_expl:
+                    new_expl = [line.replace("\n", "\\n") for line in new_expl]
+                    res = "\n~".join(new_expl)
+                    if item.config.getvalue("assertmode") == "rewrite":
+                        res = res.replace("%", "%%")
+                    return res
         return None
 
     saved_assert_hooks = util._reprcompare, util._assertion_pass
@@ -218,16 +226,40 @@ def pytest_sessionfinish(session: Session) -> None:
 def pytest_assertrepr_compare(
     config: Config, op: str, left: Any, right: Any
 ) -> list[str] | None:
+    """Return an explanation for ``left op right``.
+
+    Internally ``util.assertrepr_compare`` is a generator; we feed it
+    through ``materialize_with_truncation`` so a huge comparison
+    short-circuits at the truncation threshold without building the
+    full diff, while still returning the ``list[str] | None`` shape
+    the hook spec advertises.
+    """
     if config.pluginmanager.has_plugin("terminalreporter"):
         highlighter = config.get_terminal_writer()._highlight
     else:
         # Keep it plaintext when not using terminalrepoterer (#14377).
         highlighter = util.dummy_highlighter
-    return util.assertrepr_compare(
+    # When truncation is going to clip the explanation downstream, tell
+    # the comparison helpers to cap their pformat output at the same
+    # budget so they don't spend O(N) formatting lines we're about to
+    # drop. ``+ 3`` matches the truncator's own ``line_cap``: 2 lines
+    # for the truncation message it appends (blank + footer) plus 1
+    # for overshoot detection. ``difflib.ndiff`` over two K-line
+    # pformat outputs produces at least K output lines (more when the
+    # sides differ), and the truncator pulls at most ``trunc_lines +
+    # 3`` lines from the whole explanation, so a per-side pformat
+    # budget of ``trunc_lines + 3`` covers the worst case. With
+    # truncation disabled the cap stays ``None`` and the user gets the
+    # full diff.
+    should_truncate, trunc_lines, _ = truncate._get_truncation_parameters(config)
+    pformat_cap = trunc_lines + 3 if should_truncate and trunc_lines > 0 else None
+    lines = util.assertrepr_compare(
         op=op,
         left=left,
         right=right,
         verbose=config.get_verbosity(Config.VERBOSITY_ASSERTIONS),
         highlighter=highlighter,
         assertion_text_diff_style=util.get_assertion_text_diff_style(config),
+        pformat_cap=pformat_cap,
     )
+    return truncate.materialize_with_truncation(lines, config) or None

src/_pytest/assertion/_compare_any.py

@@ -28,54 +28,55 @@ def _compare_eq_any(
     highlighter: _HighlightFunc,
     verbose: int,
     assertion_text_diff_style: _AssertionTextDiffStyle,
-) -> list[str]:
-    explanation = []
+    pformat_cap: int | None = None,
+) -> Iterator[str]:
+    """Yield the per-line explanation for ``left == right`` (without summary).
+
+    Yields nothing when no specialised explanation applies, so consumers
+    can stream the output and bail out early (e.g. for truncation) without
+    materialising the entire diff first.
+    """
     if istext(left) and istext(right):
-        explanation = list(
-            _compare_eq_text(
-                left,
-                right,
-                highlighter,
-                verbose,
-                assertion_text_diff_style,
-            )
+        yield from _compare_eq_text(
+            left,
+            right,
+            highlighter,
+            verbose,
+            assertion_text_diff_style,
         )
     else:
         from _pytest.python_api import ApproxBase
 
         # Although the common order should be obtained == approx(...), allow both ways.
         if isinstance(right, ApproxBase):
-            explanation = right._repr_compare(left)
+            yield from right._repr_compare(left)
         elif isinstance(left, ApproxBase):
-            explanation = left._repr_compare(right)
+            yield from left._repr_compare(right)
         elif type(left) is type(right) and (
             isdatacls(left) or isattrs(left) or isnamedtuple(left)
         ):
             # Note: unlike dataclasses/attrs, namedtuples compare only the
             # field values, not the type or field names. But this branch
             # intentionally only handles the same-type case, which was often
             # used in older code bases before dataclasses/attrs were available.
-            explanation = list(
-                _compare_eq_cls(
-                    left,
-                    right,
-                    highlighter,
-                    verbose,
-                    assertion_text_diff_style,
-                )
+            yield from _compare_eq_cls(
+                left,
+                right,
+                highlighter,
+                verbose,
+                assertion_text_diff_style,
             )
         elif issequence(left) and issequence(right):
-            explanation = list(_compare_eq_sequence(left, right, highlighter, verbose))
+            yield from _compare_eq_sequence(left, right, highlighter, verbose)
         elif isset(left) and isset(right):
-            explanation = _compare_eq_set(left, right, highlighter, verbose)
+            yield from _compare_eq_set(left, right, highlighter, verbose)
         elif ismapping(left) and ismapping(right):
-            explanation = list(_compare_eq_mapping(left, right, highlighter, verbose))
+            yield from _compare_eq_mapping(left, right, highlighter, verbose)
 
         if isiterable(left) and isiterable(right):
-            expl = _compare_eq_iterable(left, right, highlighter, verbose)
-            explanation.extend(expl)
-
-    return explanation
+            yield from _compare_eq_iterable(
+                left, right, highlighter, verbose, pformat_cap
+            )
 
 
 def _compare_eq_cls(

src/_pytest/assertion/_compare_sequence.py

@@ -15,26 +15,36 @@ def _compare_eq_iterable(
     right: Iterable[object],
     highlighter: _HighlightFunc,
     verbose: int = 0,
+    pformat_cap: int | None = None,
 ) -> Iterator[str]:
     if verbose <= 0 and not running_on_ci():
         yield "Use -v to get more diff"
         return
     # dynamic import to speedup pytest
     import difflib
 
-    left_formatting = PrettyPrinter().pformat(left).splitlines()
-    right_formatting = PrettyPrinter().pformat(right).splitlines()
+    # ``pformat_cap`` is computed by the dispatcher from the
+    # truncator's ``truncation_limit_lines``: when truncation is going
+    # to drop everything past that budget anyway, we don't bother
+    # formatting more. ``None`` means no cap (``-vv`` or CI: the user
+    # wants the full diff).
+    pp = PrettyPrinter()
+    left_formatting = pp.pformat_lines(left, max_lines=pformat_cap)
+    right_formatting = pp.pformat_lines(right, max_lines=pformat_cap)
 
     yield ""
     yield "Full diff:"
     # "right" is the expected base against which we compare "left",
     # see https://github.com/pytest-dev/pytest/issues/3333
-    yield from highlighter(
-        "\n".join(
-            line.rstrip() for line in difflib.ndiff(right_formatting, left_formatting)
-        ),
-        lexer="diff",
-    ).splitlines()
+    #
+    # Yield each ndiff line through the highlighter individually so the
+    # streaming truncator can stop pulling from ``difflib.ndiff`` as
+    # soon as its budget is full. The diff lexer is line-oriented, so
+    # per-line highlighting is equivalent — it just adds a redundant
+    # ``\x1b[0m`` reset at the start of each line (invisible to the
+    # terminal).
+    for line in difflib.ndiff(right_formatting, left_formatting):
+        yield highlighter(line.rstrip(), lexer="diff")
 
 
 def _compare_eq_sequence(