diff --git a/manim/utils/iterables.py b/manim/utils/iterables.py index 196248cd73..cdd41e59d0 100644 --- a/manim/utils/iterables.py +++ b/manim/utils/iterables.py @@ -143,6 +143,7 @@ def list_difference_update(l1: Iterable[T], l2: Iterable[T]) -> list[T]: >>> list_difference_update([1, 2, 3, 4], [2, 4]) [1, 3] """ + l2 = set(l2) return [e for e in l1 if e not in l2] @@ -158,7 +159,8 @@ def list_update(l1: Iterable[T], l2: Iterable[T]) -> list[T]: >>> list_update([1, 2, 3], [2, 4, 4]) [1, 3, 2, 4, 4] """ - return [e for e in l1 if e not in l2] + list(l2) + l2 = list(l2) + return list_difference_update(l1, l2) + l2 @overload diff --git a/tests/module/utils/test_iterables.py b/tests/module/utils/test_iterables.py new file mode 100644 index 0000000000..6f52d7441b --- /dev/null +++ b/tests/module/utils/test_iterables.py @@ -0,0 +1,35 @@ +import pytest + +from manim.utils.iterables import list_difference_update, list_update + + +def test_list_difference_update_removes_matching_items(): + assert list_difference_update([1, 2, 3, 4], [2, 4]) == [1, 3] + + +def test_list_difference_update_preserves_duplicates_and_order(): + assert list_difference_update([3, 1, 3, 2, 1], [1]) == [3, 3, 2] + + +@pytest.mark.parametrize( + ("l1", "l2", "expected"), + [([1, 2, 3, 4], [2, 4], [1, 3]), ([], [1, 2], [])], +) +def test_list_difference_update_removes_elements(l1, l2, expected): + assert list_difference_update(l1, l2) == expected + + +def test_list_difference_update_preserves_l1_order_and_duplicates(): + assert list_difference_update([3, 1, 3, 2, 1], [1]) == [3, 3, 2] + + +@pytest.mark.parametrize( + ("l1", "l2", "expected"), + [([1, 2, 3], [2, 4], [1, 3, 2, 4]), ([], [1, 2], [1, 2])], +) +def test_list_update_removes_overlap_and_appends_l2(l1, l2, expected): + assert list_update(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]