pytest/src/_pytest/assertion/_compare_set.py at e35f1d97482ee22abe31eea2ebbd6421db2bb98e · pytest-dev/pytest · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
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,
}