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

[nikolajmunk]

Overview: What does this pull request change?

This addresses a small regression introduced by #4939, where an unhashable element in l1 and/or l2 would throw a TypeError. This PR adds a key function to the signatures of list_difference_update and list_update. When provided, this key function is used to determine uniqueness/equality for all elements in l1 and l2. This behavior is similar to that of sorted and min/max, which all provide a similar key function parameter.

The list_difference function uses its current implementation if key has its default value of None, since it is slightly faster to hash directly than to pass everything through the key function. Otherwise, performance depends entirely on the complexity of the key function.

While the Manim library does not currently use either list_difference_update nor list_update with unhashable elements anywhere, intended usage is as follows: a user has, say, a list of deeply-nested dict objects and they want to remove certain dicts (or dicts with the same contents) from that list. iterables.py contains the function hash_obj which produces a hash from mutable collections, so the user writes

new_list = list_difference_update(list_of_dicts, dicts_to_remove, key=hash_obj)

to achieve their goal.

However, the flexibility afforded by key also permits more complex behavior. Here are some examples of what is possible by passing various key functions:

>>> list_difference_update(["a", "b", "A", "C"], ["A", "D"], key=str.lower)
['b', 'C']

>>> list_difference_update([1, [1], 2, 3, [3]], [[1], 3], key=str)
[1, 2, [3]]

>>> square = Square(stroke_color=RED)
>>> circle = Circle(stroke_color=RED)
>>> triangle = Triangle(stroke_color=BLUE)
>>> star = Star(stroke_color=RED)
>>> list_difference_update([square, circle, triangle], [star], key=lambda m: m.stroke_color)
[Triangle]

I've added tests for the new behavior as well as typing overloads expressing the intention that if list_difference_update and list_update are used without a key function, then l1 and l2 should be iterables of hashable elements.

Happy to hear any thoughts :)

Reviewer Checklist

• The PR title is descriptive enough for the changelog, and the PR is labeled correctly
• If applicable: newly added non-private functions and classes have a docstring including a short summary and a PARAMETERS section
• If applicable: newly added functions and classes are tested

[GniLudio]

Optimizing hashable types while still allowing unhashable types would be nice.
Maybe with something like this?

def list_difference_update(l1: Iterable[T], l2: Iterable[T]) -> list[T]:
    if isinstance(l2, Iterator):
        l2 = tuple(l2)
    with contextlib.suppress(TypeError):
        l2 = set(l2)
    return [e for e in l1 if e not in l2]

[nikolajmunk]

Yep, that works for cases when l2 has unhashables. But if, say, l1 has an unhashable element while l2 doesn't, then the list comprehension will try to look up an unhashable in a set and error out.

[GniLudio]

Yep, that works for cases when l2 has unhashables. But if, say, l1 has an unhashable element while l2 doesn't, then the list comprehension will try to look up an unhashable in a set and error out.

Not really liking this solution:

def list_difference_update(l1, l2):
    if isinstance(l2, Iterator):
        l2 = tuple(l2)
    
    with contextlib.suppress(TypeError):
        l2_set = set(l2)
        
    def is_in_l2(e) -> bool:
        try:
            return e in l2_set
        except TypeError:
            return e in l2
    return [e for e in l1 if not is_in_l2(e)]

Here a small test

from typing import Iterator
import contextlib

def list_difference_update(l1, l2):
    if isinstance(l2, Iterator):
        l2 = tuple(l2)
    
    with contextlib.suppress(TypeError):
        l2_set = set(l2)
        
    def is_in_l2(e) -> bool:
        try:
            return e in l2_set
        except TypeError:
            return e in l2
    return [e for e in l1 if not is_in_l2(e)]
    
    
class MyClass:
    def __init__(self, name):
        self.name = name
        
    def __repr__(self):
        return self.name
        
    __hash__ = None

a = [MyClass("a"), MyClass("b"), MyClass("c")]
b = [1, 2, 3]
c = list_difference_update(a, b)
print(c)

[nikolajmunk]

Not really liking this solution:

Haha yes, I have a few solutions that look similar. It's possible to do this in a nice-looking way, but all my attempts have resulted in an algorithm that is markedly slower when all input is hashable, which it will be the vast majority of the time. Catching the edge case is simply very expensive.

[nikolajmunk]

Okay, two things now that I've had a change to think about it.

1. I think it would be entirely fine to just not support unhashable items in the input. If people want to do a difference update on, say, a list of sets, they'll have to implement their own function.
2. If we really want the library to contain a difference function for unhashables, could we add an optional key parameter to the function, similar to sort or max? Implementation could look like this:

def list_difference_update(
        l1: Iterable[T], l2: Iterable[U], key: Callable[[T | U], Hashable] | None = None
    ) -> list[T]:
    key = hash if key is None else key
    l2_keys = {key(e) for e in l2}
    return [e for e in l1 if key(e) not in l2_keys]

Using this implementation, all the following assertions pass:

from manim.utils.iterables import list_difference_update, hash_obj
assert list_difference_update([1, 2, 3, 4], [2, 3]) == [1, 4]
assert list_difference_update([[1], [2], [3], [4]], [[2], [3]], key=hash_obj) == [[1], [4]]
assert list_difference_update([1, 2, [3], 4], [2, 3], key=hash_obj) == [1, [3], 4]
assert list_difference_update([1, 2, [3], 4], [2, [3]], key=hash_obj) == [1, 4]
assert list_difference_update([[1], [2], [3]], [2, 3], key=hash_obj) == [[1], [2], [3]]

This of course also broadens the possible uses of the function, which we may or may not want. For example, it would be possible to use it to filter out any items in l1 with the same length as anything in l2:

>>> list_difference_update(["a", "b", "ab", "abc", [1, 2, 3], [1, 2]], [[0], "000"], key=len)
['ab', [1, 2]]

I haven't profiled this at all yet so maybe the performance cost is simply too high (for "well-behaved" input) to justify doing this, but at least it's possible and doesn't require any manual "is this hashable" checking on the function's part.

Thoughts?

[nikolajmunk]

I banged together a quick benchmark, and it seems like it would work pretty well to mix the above suggestion with the existing implementation:

def list_difference_update(
    l1: Iterable[T],
    l2: Iterable[U],
    key: Callable[[T | U], Hashable] | None = None,
) -> list[T]:
    if key in (None, hash):
        l2 = set(l2)
        return [e for e in l1 if e not in l2]
    else:
        l2 = set(map(key, l2))
        return [e for e in l1 if key(e) not in l2]

This yields the following per-element running times (here, a hashable element is just an object() while an unhashable object is [object()]):


This keeps the default hashable path fast! (And perhaps hash_obj_with_cache is a little misleading since if we have a function yielding unique results for use in the memoization table, then we could just use that as our hash function in the first place!)

[nikolajmunk]

I went ahead and pushed the above implementation. It needs a bunch of overloads to capture the intended behavior when key isn't provided, but everything seems to work.

This probably needs some tests, I'll see what I can do.

[nikolajmunk]

Added some new tests for the key function and for handling of various different iterable types, particularly consumable iterables and non-Sequences.

This is a lot of testing for two functions, do people think it's overkill?