1 from __future__ import annotations

2

3 from muutils.group_equiv import group_by_equivalence

4

5

6 def test_group_by_equivalence_simple():

7 # Test with integers and simple equality

8 assert group_by_equivalence([1, 1, 2, 2, 3, 3], lambda x, y: x == y) == [

9 [1, 1],

10 [2, 2],

11 [3, 3],

12 ]

13 assert group_by_equivalence([1, 1, 2, 4, 5, 5], lambda x, y: x == y) == [

14 [1, 1],

15 [2],

16 [4],

17 [5, 5],

18 ]

19

20 # Test with strings and simple equality

21 assert group_by_equivalence(

22 ["a", "a", "b", "b", "c", "c"], lambda x, y: x == y

23 ) == [["a", "a"], ["b", "b"], ["c", "c"]]

24

25

26 def _non_transitive_relation(x: int, y: int) -> bool:

27 return abs(x - y) <= 1

28

29

30 def test_group_by_equivalence_non_transitive():

31 assert group_by_equivalence(list(range(5)), _non_transitive_relation) == [

32 [0, 1, 2, 3, 4]

33 ]

34

35 assert group_by_equivalence([0, 1, 10, 11], _non_transitive_relation) == [

36 [0, 1],

37 [10, 11],

38 ]

39

40 assert group_by_equivalence(

41 [0, 1, 5, 4, 10, 11, 3, 2], _non_transitive_relation

42 ) == [[10, 11], [0, 1, 5, 4, 3, 2]]

43

44 assert group_by_equivalence(

45 [0, 1, 5, 4, 10, 11, 3, 2], _non_transitive_relation

46 ) == [[10, 11], [0, 1, 5, 4, 3, 2]]