[refactor] Create the SET_COMPARISON_FUNCTIONS once

Instead of creating it at runtime all the time

Author: Pierre-Sassoulas

src/_pytest/assertion/_compare_set.py

@@ -0,0 +1,97 @@
+from __future__ import annotations
+
+from collections.abc import Callable
+from collections.abc import Set as AbstractSet
+from typing import Any
+
+from _pytest._io.saferepr import saferepr
+from _pytest.assertion._typing import _HighlightFunc
+
+
+def _set_one_sided_diff(
+    posn: str,
+    set1: AbstractSet[Any],
+    set2: AbstractSet[Any],
+    highlighter: _HighlightFunc,
+) -> list[str]:
+    explanation = []
+    diff = set1 - set2
+    if diff:
+        explanation.append(f"Extra items in the {posn} set:")
+        for item in diff:
+            explanation.append(highlighter(saferepr(item)))
+    return explanation
+
+
+def _compare_eq_set(
+    left: AbstractSet[Any],
+    right: AbstractSet[Any],
+    highlighter: _HighlightFunc,
+    verbose: int = 0,
+) -> list[str]:
+    explanation = []
+    explanation.extend(_set_one_sided_diff("left", left, right, highlighter))
+    explanation.extend(_set_one_sided_diff("right", right, left, highlighter))
+    return explanation
+
+
+def _compare_gt_set(
+    left: AbstractSet[Any],
+    right: AbstractSet[Any],
+    highlighter: _HighlightFunc,
+    verbose: int = 0,
+) -> list[str]:
+    explanation = _compare_gte_set(left, right, highlighter)
+    if not explanation:
+        return ["Both sets are equal"]
+    return explanation
+
+
+def _compare_lt_set(
+    left: AbstractSet[Any],
+    right: AbstractSet[Any],
+    highlighter: _HighlightFunc,
+    verbose: int = 0,
+) -> list[str]:
+    explanation = _compare_lte_set(left, right, highlighter)
+    if not explanation:
+        return ["Both sets are equal"]
+    return explanation
+
+
+def _compare_gte_set(
+    left: AbstractSet[Any],
+    right: AbstractSet[Any],
+    highlighter: _HighlightFunc,
+    verbose: int = 0,
+) -> list[str]:
+    return _set_one_sided_diff("right", right, left, highlighter)
+
+
+def _compare_lte_set(
+    left: AbstractSet[Any],
+    right: AbstractSet[Any],
+    highlighter: _HighlightFunc,
+    verbose: int = 0,
+) -> list[str]:
+    return _set_one_sided_diff("left", left, right, highlighter)
+
+
+SetComparisonFunction = dict[
+    str,
+    Callable[
+        [AbstractSet[Any], AbstractSet[Any], _HighlightFunc, int],
+        list[str],
+    ],
+]
+
+SET_COMPARISON_FUNCTIONS: SetComparisonFunction = {
+    # == can't be done here without a prior refactor because there's an additional
+    # explanation for iterable in _compare_eq_any
+    # "==": _compare_eq_set,
+    "!=": lambda *a, **kw: ["Both sets are equal"],
+    ">=": _compare_gte_set,
+    "<=": _compare_lte_set,
+    ">": _compare_gt_set,
+    "<": _compare_lt_set,
+}

src/_pytest/assertion/_typing.py

@@ -0,0 +1,9 @@
+from __future__ import annotations
+
+from typing import Literal
+from typing import Protocol
+
+
+class _HighlightFunc(Protocol):  # noqa: PYI046
+    def __call__(self, source: str, lexer: Literal["diff", "python"] = "python") -> str:
+        """Apply highlighting to the given source."""

src/_pytest/assertion/util.py

@@ -8,19 +8,20 @@
 from collections.abc import Iterable
 from collections.abc import Mapping
 from collections.abc import Sequence
-from collections.abc import Set as AbstractSet
 import os
 import pprint
 from typing import Any
 from typing import Literal
-from typing import Protocol
 from unicodedata import normalize
 
 from _pytest import outcomes
 import _pytest._code
 from _pytest._io.pprint import PrettyPrinter
 from _pytest._io.saferepr import saferepr
 from _pytest._io.saferepr import saferepr_unlimited
+from _pytest.assertion._compare_set import _compare_eq_set
+from _pytest.assertion._compare_set import SET_COMPARISON_FUNCTIONS
+from _pytest.assertion._typing import _HighlightFunc
 from _pytest.config import Config
 
 
@@ -38,20 +39,6 @@
 _config: Config | None = None
 
 
-class _HighlightFunc(Protocol):
-    def __call__(self, source: str, lexer: Literal["diff", "python"] = "python") -> str:
-        """Apply highlighting to the given source."""
-
-
-CompareSetFunction = dict[
-    str,
-    Callable[
-        [AbstractSet[Any], AbstractSet[Any], _HighlightFunc, int],
-        list[str],
-    ],
-]
-
-
 def dummy_highlighter(source: str, lexer: Literal["diff", "python"] = "python") -> str:
     """Dummy highlighter that returns the text unprocessed.
 
@@ -225,14 +212,9 @@ def assertrepr_compare(
                 "!=" | ">=" | "<=" | ">" | "<",
                 set() | frozenset(),
             ):
-                set_compare_func: CompareSetFunction = {
-                    "!=": lambda *a, **kw: ["Both sets are equal"],
-                    ">=": _compare_gte_set,
-                    "<=": _compare_lte_set,
-                    ">": _compare_gt_set,
-                    "<": _compare_lt_set,
-                }
-                explanation = set_compare_func[op](left, right, highlighter, verbose)
+                explanation = SET_COMPARISON_FUNCTIONS[op](
+                    left, right, highlighter, verbose
+                )
             case _:
                 explanation = None
     except outcomes.Exit:
@@ -433,75 +415,6 @@ def _compare_eq_sequence(
     return explanation
 
 
-def _compare_eq_set(
-    left: AbstractSet[Any],
-    right: AbstractSet[Any],
-    highlighter: _HighlightFunc,
-    verbose: int = 0,
-) -> list[str]:
-    explanation = []
-    explanation.extend(_set_one_sided_diff("left", left, right, highlighter))
-    explanation.extend(_set_one_sided_diff("right", right, left, highlighter))
-    return explanation
-
-
-def _compare_gt_set(
-    left: AbstractSet[Any],
-    right: AbstractSet[Any],
-    highlighter: _HighlightFunc,
-    verbose: int = 0,
-) -> list[str]:
-    explanation = _compare_gte_set(left, right, highlighter)
-    if not explanation:
-        return ["Both sets are equal"]
-    return explanation
-
-
-def _compare_lt_set(
-    left: AbstractSet[Any],
-    right: AbstractSet[Any],
-    highlighter: _HighlightFunc,
-    verbose: int = 0,
-) -> list[str]:
-    explanation = _compare_lte_set(left, right, highlighter)
-    if not explanation:
-        return ["Both sets are equal"]
-    return explanation
-
-
-def _compare_gte_set(
-    left: AbstractSet[Any],
-    right: AbstractSet[Any],
-    highlighter: _HighlightFunc,
-    verbose: int = 0,
-) -> list[str]:
-    return _set_one_sided_diff("right", right, left, highlighter)
-
-
-def _compare_lte_set(
-    left: AbstractSet[Any],
-    right: AbstractSet[Any],
-    highlighter: _HighlightFunc,
-    verbose: int = 0,
-) -> list[str]:
-    return _set_one_sided_diff("left", left, right, highlighter)
-
-
-def _set_one_sided_diff(
-    posn: str,
-    set1: AbstractSet[Any],
-    set2: AbstractSet[Any],
-    highlighter: _HighlightFunc,
-) -> list[str]:
-    explanation = []
-    diff = set1 - set2
-    if diff:
-        explanation.append(f"Extra items in the {posn} set:")
-        for item in diff:
-            explanation.append(highlighter(saferepr(item)))
-    return explanation
-
-
 def _compare_eq_dict(
     left: Mapping[Any, Any],
     right: Mapping[Any, Any],