Repository: GridTools/gt4py
Pull request #2062 (2025-06-04 07:41:48)

feat[next]: improve eve caching and hashing utilities
Enhance custom hashing utilities with a `CustomMapping` class and performance improvements in `HashableBy`. Additionally, new unit tests have been added and some of these classes are moved around the `eve.utils` module for more coherent organization.


Patch:
diff --git a/src/gt4py/eve/utils.py b/src/gt4py/eve/utils.py
index 06acb619..e4bfe84d 100644
--- a/src/gt4py/eve/utils.py
+++ b/src/gt4py/eve/utils.py
@@ -295,53 +295,109 @@ class CustomDefaultDictBase(collections.UserDict[_K, _V]):
             return value
 
 
-class fluid_partial(functools.partial):
-    """Create a `functools.partial` with support for multiple applications calling `.partial()`."""
+class CustomMapping(collections.abc.MutableMapping[_K, _V]):
+    """
+    A custom mapping class that uses a key function to map keys to values.
 
-    def partial(self, *args: Any, **kwargs: Any) -> fluid_partial:
-        return fluid_partial(self, *args, **kwargs)
+    This class allows for custom key functions to be used for indexing and
+    retrieving values, while maintaining a mapping-like interface.
 
+    Examples:
+        >>> mapping = CustomMapping(lambda x: hash(repr(x)))
+        >>> mapping[[1, 2]] = "first"
+        >>> mapping[[1, 2]]
+        'first'
+        >>> mapping[{1, 2}] = "second"
+        >>> mapping[{1, 2}]
+        'second'
+        >>> len(mapping)
+        2
+    """
 
-@overload
-def with_fluid_partial(
-    func: Literal[None] = None, *args: Any, **kwargs: Any
-) -> Callable[[Callable[_P, _T]], Callable[_P, _T]]: ...
+    __slots__ = ("key_func", "key_map", "value_map")
 
+    key_func: Callable[[_K], int]
+    key_map: dict[int, _K]
+    value_map: dict[int, _V]
 
-@overload
-def with_fluid_partial(  # redefinition of unused function
-    func: Callable[_P, _T], *args: Any, **kwargs: Any
-) -> Callable[_P, _T]: ...
+    def __init__(self, key: Callable[[_K], int]) -> None:
+        self.key_func = key
+        self.key_map = {}
+        self.value_map = {}
 
+    def __getitem__(self, key: _K) -> _V:
+        return self.value_map[self.key_func(key)]
+
+    def __setitem__(self, key: _K, value: _V) -> None:
+        data_key = self.key_func(key)
+        self.key_map[data_key] = key
+        self.value_map[data_key] = value
+
+    def __delitem__(self, key: _K) -> None:
+        custom_key = self.key_func(key)
+        del self.key_map[custom_key]
+        del self.value_map[custom_key]
+
+    def __len__(self) -> int:
+        return len(self.key_map)
+
+    def __iter__(self) -> Iterator[_K]:
+        for custom_key in self.key_map:
+            yield self.key_map[custom_key]
+
+    def __repr__(self) -> str:
+        return (
+            "{"
+            + ", ".join(f"{self.key_map[k]!r}: {self.value_map[k]!r}" for k in self.key_map)
+            + "}"
+        )
 
-def with_fluid_partial(  # redefinition of unused function
-    func: Optional[Callable[..., Any]] = None, *args: Any, **kwargs: Any
-) -> Union[Callable[..., Any], Callable[[Callable[..., Any]], Callable[..., Any]]]:
-    """Add a `partial` attribute to the decorated function.
 
-    The `partial` attribute is a function that behaves like `functools.partial`,
-    but also supports partial application of the decorated function. It can be
-    used both as a bare or a parameterized decorator.
+class HashableBy(Generic[_T]):
+    __slots__ = ("hashed_value", "value")
 
-    Arguments:
-        func: The function to decorate.
+    def __init__(self, hash_func: Callable[[_T], int], value: _T) -> None:
+        self.value = value
+        self.hashed_value = hash_func(value)
 
-    Returns:
-        Returns the decorated function with an extra `.partial()` attribute.
+    def __hash__(self) -> int:
+        return self.hashed_value
 
-    Example:
-        >>> @with_fluid_partial
-        ... def add(a, b):
-        ...     return a + b
-        >>> add.partial(1)(2)
-        3
+    def __eq__(self, other: Any) -> bool:
+        assert isinstance(other, HashableBy)
+        return self.value is other.value or self.value == other.value
+
+    def __str__(self) -> str:
+        return (
+            f"{self.__class__.__name__}(value={self.value!r}, hashed_value={self.hashed_value!r})"
+        )
+
+
+@overload
+def hashable_by(func: Callable[[_T], int], value: _T) -> HashableBy[_T]: ...
+
+
+@overload
+def hashable_by(
+    func: Callable[[_T], int], value: NothingType = NOTHING
+) -> functools.partial[HashableBy[_T]]: ...
+
+
+def hashable_by(
+    func: Callable[[_T], int], value: _T | NothingType = NOTHING
+) -> HashableBy[_T] | functools.partial[HashableBy[_T]]:
+    """
+    Creates a wrapper that uses `func` to hash the value passed to it.
     """
+    return (
+        HashableBy(func, value)  # type: ignore[arg-type] # checked in condition
+        if value is not NOTHING
+        else functools.partial(hashable_by, func)
+    )
 
-    def _decorator(func: Callable[..., Any]) -> Callable[..., Any]:
-        func.partial = fluid_partial(functools.partial, func, *args, **kwargs)  # type: ignore[attr-defined]  # add attribute
-        return func
 
-    return _decorator(func) if func is not None else _decorator
+hashable_by_id = hashable_by(id)
+cached_hash = hashable_by(hash)
 
 
 @overload
@@ -351,12 +407,12 @@ def optional_lru_cache(
 
 
 @overload
-def optional_lru_cache(  # redefinition of unused function
+def optional_lru_cache(
     func: Callable[_P, _T], *, maxsize: Optional[int] = 128, typed: bool = False
 ) -> Callable[_P, _T]: ...
 
 
-def optional_lru_cache(  # redefinition of unused function
+def optional_lru_cache(
     func: Optional[Callable[_P, _T]] = None, *, maxsize: Optional[int] = 128, typed: bool = False
 ) -> Union[Callable[_P, _T], Callable[[Callable[_P, _T]], Callable[_P, _T]]]:
     """Wrap :func:`functools.lru_cache` to fall back to the original function if arguments are not hashable.
@@ -401,46 +457,6 @@ def optional_lru_cache(  # redefinition of unused function
     return _decorator(func) if func is not None else _decorator
 
 
-class HashableBy(Generic[_T]):
-    __slots__ = ["func", "value"]
-
-    def __init__(self, func: Callable[[_T], int], value: _T) -> None:
-        self.func = func
-        self.value = value
-
-    def __hash__(self) -> int:
-        return self.func(self.value)
-
-    def __eq__(self, other: Any) -> bool:
-        return self.value is other.value and self.func is self.func
-
-
-@overload
-def hashable_by(func: Callable[[_T], int], value: _T) -> HashableBy[_T]: ...
-
-
-@overload
-def hashable_by(
-    func: Callable[[_T], int], value: NothingType = NOTHING
-) -> functools.partial[HashableBy[_T]]: ...
-
-
-def hashable_by(
-    func: Callable[[_T], int], value: _T | NothingType = NOTHING
-) -> HashableBy[_T] | functools.partial[HashableBy[_T]]:
-    """
-    Creates a wrapper that uses `func` to hash the value passed to it.
-    """
-    return (
-        HashableBy(func, value)  # type: ignore[arg-type] # checked in condition
-        if value is not NOTHING
-        else functools.partial(hashable_by, func)
-    )
-
-
-hashable_by_id = hashable_by(id)
-
-
 # TODO(egparedes): it would be more efficient to implement the caching logic
 # here instead of relying on `functools.lru_cache` and wrapping/unwrapping the
 # arguments.
@@ -497,6 +513,53 @@ def lru_cache(
     return _decorator(func) if func is not None else _decorator
 
 
+class fluid_partial(functools.partial):
+    """Create a `functools.partial` with support for multiple applications calling `.partial()`."""
+
+    def partial(self, *args: Any, **kwargs: Any) -> fluid_partial:
+        return fluid_partial(self, *args, **kwargs)
+
+
+@overload
+def with_fluid_partial(
+    func: Literal[None] = None, *args: Any, **kwargs: Any
+) -> Callable[[Callable[_P, _T]], Callable[_P, _T]]: ...
+
+
+@overload
+def with_fluid_partial(func: Callable[_P, _T], *args: Any, **kwargs: Any) -> Callable[_P, _T]: ...
+
+
+def with_fluid_partial(
+    func: Optional[Callable[..., Any]] = None, *args: Any, **kwargs: Any
+) -> Union[Callable[..., Any], Callable[[Callable[..., Any]], Callable[..., Any]]]:
+    """Add a `partial` attribute to the decorated function.
+
+    The `partial` attribute is a function that behaves like `functools.partial`,
+    but also supports partial application of the decorated function. It can be
+    used both as a bare or a parameterized decorator.
+
+    Arguments:
+        func: The function to decorate.
+
+    Returns:
+        Returns the decorated function with an extra `.partial()` attribute.
+
+    Example:
+        >>> @with_fluid_partial
+        ... def add(a, b):
+        ...     return a + b
+        >>> add.partial(1)(2)
+        3
+    """
+
+    def _decorator(func: Callable[..., Any]) -> Callable[..., Any]:
+        func.partial = fluid_partial(functools.partial, func, *args, **kwargs)  # type: ignore[attr-defined]  # add attribute
+        return func
+
+    return _decorator(func) if func is not None else _decorator
+
+
 def register_subclasses(*subclasses: Type) -> Callable[[Type], Type]:
     """Class decorator to automatically register virtual subclasses.
 


Test patch:
diff --git a/tests/eve_tests/unit_tests/test_utils.py b/tests/eve_tests/unit_tests/test_utils.py
index 448624c1..671d1672 100644
--- a/tests/eve_tests/unit_tests/test_utils.py
+++ b/tests/eve_tests/unit_tests/test_utils.py
@@ -151,19 +151,74 @@ def unique_data_items(request):
     ]
 
 
-def test_fluid_partial():
-    from gt4py.eve.utils import fluid_partial
+def test_custom_default_dict_base():
+    from gt4py.eve.utils import CustomDefaultDictBase
 
-    def func(a, b, c):
-        return a + b + c
+    class TestDefaultDict(CustomDefaultDictBase):
+        def value_factory(self, key):
+            return key * 2
 
-    fp1 = fluid_partial(func, 1)
-    fp2 = fp1.partial(2)
-    fp3 = fp2.partial(3)
+    d = TestDefaultDict()
 
-    assert fp1(2, 3) == 6
-    assert fp2(3) == 6
-    assert fp3() == 6
+    # Test default value creation
+    assert d[5] == 10
+    assert d[3] == 6
+
+    # Test override
+    d[5] = 100
+    assert d[5] == 100
+
+    # Test length
+    assert len(d) == 2
+
+    # Test iteration
+    keys = list(d.keys())
+    assert 5 in keys and 3 in keys
+
+
+def test_custom_mapping():
+    from gt4py.eve.utils import CustomMapping
+
+    keys = [[1, 2], {"foo", "bar"}]
+    values = ["value1", "value2"]
+
+    mapping = CustomMapping(lambda x: hash(repr(x)))
+    for key, value in zip(keys, values):
+        mapping[key] = value
+
+    assert len(mapping) == 2
+    assert all(mapping[key] == value for key, value in zip(keys, values))
+
+    # Test iteration
+    assert keys == list(mapping)
+
+    # Test deletion
+    del mapping[keys[0]]
+    assert len(mapping) == 1
+    assert keys[-1] in mapping
+
+    with pytest.raises(KeyError):
+        _ = mapping[[3, 4]]
+
+    # Test with different key function
+    id_mapping = CustomMapping(id)
+    obj1 = [1, 2, 3]
+    obj2 = [1, 2, 3]
+
+    id_mapping[obj1] = "obj1"
+    id_mapping[obj2] = "obj2"
+
+    assert id_mapping[obj1] == "obj1"
+    assert id_mapping[obj2] == "obj2"
+    assert len(id_mapping) == 2
+
+
+def test_HashableBy():
+    from gt4py.eve.utils import HashableBy
+
+    assert hash(HashableBy(id, 345)) == id(345)
+    assert "value=345" in str(HashableBy(lambda x: "FOO", 345))
+    assert "hashed_value='FOO'" in str(HashableBy(lambda x: "FOO", 345))
 
 
 def test_hashable_by():
@@ -184,6 +239,52 @@ def test_hashable_by_id():
     assert hash(hashable_by_id(testee)) == id(testee)
 
 
+def test_cached_hash():
+    from gt4py.eve.utils import cached_hash
+
+    testee = (1, 2)
+
+    assert hash(cached_hash(testee)) == hash(testee)
+
+
+def test_lru_cache_key_id_called_once():
+    from gt4py.eve.utils import lru_cache
+
+    call_count = 0
+
+    def func(x):
+        nonlocal call_count
+        call_count += 1
+        return x
+
+    cached = lru_cache(func, key=id)
+
+    assert cached.__wrapped__ == func
+
+    obj = object()
+    assert cached(obj) is obj
+    assert cached(obj) is obj
+    assert call_count == 1
+
+    assert cached.cache_info().hits == 1
+    assert cached.cache_info().misses == 1
+
+
+def test_fluid_partial():
+    from gt4py.eve.utils import fluid_partial
+
+    def func(a, b, c):
+        return a + b + c
+
+    fp1 = fluid_partial(func, 1)
+    fp2 = fp1.partial(2)
+    fp3 = fp2.partial(3)
+
+    assert fp1(2, 3) == 6
+    assert fp2(3) == 6
+    assert fp3() == 6
+
+
 def test_noninstantiable_class():
     @eve.utils.noninstantiable
     class NonInstantiableClass(eve.datamodels.DataModel):
@@ -369,26 +470,3 @@ def test_xenumerate():
     from gt4py.eve.utils import xenumerate
 
     assert list(xenumerate(string.ascii_letters[:3])) == [(0, "a"), (1, "b"), (2, "c")]
-
-
-def test_lru_cache_key_id_called_once():
-    from gt4py.eve.utils import lru_cache
-
-    call_count = 0
-
-    def func(x):
-        nonlocal call_count
-        call_count += 1
-        return x
-
-    cached = lru_cache(func, key=id)
-
-    assert cached.__wrapped__ == func
-
-    obj = object()
-    assert cached(obj) is obj
-    assert cached(obj) is obj
-    assert call_count == 1
-
-    assert cached.cache_info().hits == 1
-    assert cached.cache_info().misses == 1