diff --git a/manim/utils/iterables.py b/manim/utils/iterables.py index 66708353a5..bb6ce89cff 100644 --- a/manim/utils/iterables.py +++ b/manim/utils/iterables.py @@ -36,6 +36,7 @@ U = TypeVar("U") F = TypeVar("F", np.float64, np.int_) H = TypeVar("H", bound=Hashable) +J = TypeVar("J", bound=Hashable) if TYPE_CHECKING: @@ -133,24 +134,85 @@ def concatenate_lists(*list_of_lists: Iterable[T]) -> list[T]: return [item for lst in list_of_lists for item in lst] -def list_difference_update(l1: Iterable[T], l2: Iterable[U]) -> list[T]: +@overload +def list_difference_update( + l1: Iterable[H], l2: Iterable[J], *, key: None = None +) -> list[H]: ... +@overload +def list_difference_update( + l1: Iterable[T], + l2: Iterable[U], + *, + key: Callable[[T | U], Hashable], +) -> list[T]: ... +def list_difference_update( + l1: Iterable[T], + l2: Iterable[U], + *, + key: Callable[[T | U], Hashable] | None = None, +) -> list[T]: """Returns a list containing all the elements of l1 not in l2. + Parameters + ---------- + l1 + The first iterable. + l2 + The second iterable. + key + A key function which provides a value used to determine uniqueness. The default + value of ``None`` means that the elements' own hash values will be used. + Examples -------- .. code-block:: pycon - >>> list_difference_update([1, 2, 3, 4], [2, 4]) - [1, 3] + >>> list_difference_update([1, 2, 3, 4, 1], [2, 4]) + [1, 3, 1] + >>> list_difference_update(["a", "b", "A", "C"], ["A", "D"], key=str.lower) + ['b', 'C'] """ - l2 = set(l2) - return [e for e in l1 if e not in l2] + if key in (None, hash): + if not isinstance(l2, (set, dict, frozenset)): + # l2 is not a set-like object, so convert it to a set for faster lookups + l2 = set(l2) + return [e for e in l1 if e not in l2] + # Use provided key function to determine uniqueness + l2_keys = set(map(key, l2)) + return [e for e in l1 if key(e) not in l2_keys] -def list_update(l1: Iterable[T], l2: Iterable[U]) -> list[T | U]: - """Used instead of ``set.update()`` to maintain order, - making sure duplicates are removed from l1, not l2. - Removes overlap of l1 and l2 and then concatenates l2 unchanged. + +@overload +def list_update( + l1: Iterable[H], l2: Iterable[J], *, key: None = None +) -> list[H | J]: ... +@overload +def list_update( + l1: Iterable[T], + l2: Iterable[U], + *, + key: Callable[[T | U], Hashable], +) -> list[T | U]: ... +def list_update( + l1: Iterable[T], + l2: Iterable[U], + *, + key: Callable[[T | U], Hashable] | None = None, +) -> list[T | U]: + """Used instead of ``set.update()`` to maintain order, making sure duplicates are + removed from l1, not l2. + Removes overlap of l1 and l2 and then concatenates l2 unchanged. + + Parameters + ---------- + l1 + The first iterable. + l2 + The second iterable. + key + A key function which provides a value used to determine uniqueness. The default + value of ``None`` means that the elements' own hash values will be used. Examples -------- @@ -158,9 +220,12 @@ def list_update(l1: Iterable[T], l2: Iterable[U]) -> list[T | U]: >>> list_update([1, 2, 3], [2, 4, 4]) [1, 3, 2, 4, 4] + >>> list_update(["a", "b", "c", "A", "B", "C"], ["A", "b", "D"], key=str.lower) + ['c', 'C', 'A', 'b', 'D'] """ - l2 = list(l2) - return list_difference_update(l1, l2) + cast(list[T | U], l2) + if not isinstance(l2, list): + l2 = list(l2) + return list_difference_update(l1, l2, key=key) + cast(list[T | U], l2) @overload diff --git a/tests/module/utils/test_iterables.py b/tests/module/utils/test_iterables.py index 6f52d7441b..498ff38920 100644 --- a/tests/module/utils/test_iterables.py +++ b/tests/module/utils/test_iterables.py @@ -23,6 +23,21 @@ def test_list_difference_update_preserves_l1_order_and_duplicates(): assert list_difference_update([3, 1, 3, 2, 1], [1]) == [3, 3, 2] +def test_list_difference_update_with_key_none_uses_default_equality(): + assert list_difference_update([1, 2, 3], [2], key=None) == [1, 3] + + +@pytest.mark.parametrize( + ("l1", "l2", "key", "expected"), + [ + (["a", "b", "A", "C"], ["A", "D"], str.lower, ["b", "C"]), + ([1, [1], 2, 3, [3]], [[1], 3], str, [1, 2, [3]]), + ], +) +def test_list_difference_update_uses_key_function(l1, l2, key, expected): + assert list_difference_update(l1, l2, key=key) == expected + + @pytest.mark.parametrize( ("l1", "l2", "expected"), [([1, 2, 3], [2, 4], [1, 3, 2, 4]), ([], [1, 2], [1, 2])], @@ -33,3 +48,18 @@ def test_list_update_removes_overlap_and_appends_l2(l1, l2, expected): def test_list_update_preserves_duplicates_in_l2(): assert list_update([1, 2, 3], [2, 4, 4]) == [1, 3, 2, 4, 4] + + +def test_list_update_with_key_none_uses_default_equality(): + assert list_update([1, 2, 3], [2, 4], key=None) == [1, 3, 2, 4] + + +@pytest.mark.parametrize( + ("l1", "l2", "key", "expected"), + [ + (["a", "b", "A", "C"], ["A", "B", "D"], str.lower, ["C", "A", "B", "D"]), + ([1, [1], 2, 3, [3]], [[1], 3], str, [1, 2, [3], [1], 3]), + ], +) +def test_list_update_uses_key_function(l1, l2, key, expected): + assert list_update(l1, l2, key=key) == expected