Add key function to handle unhashable element input in list_update and list_difference_update

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,34 +134,98 @@
     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
     --------
     .. code-block:: pycon
 
         >>> 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

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