Support grouping by keys in LatestValueCache

`LatestValueCache` now takes an optional `key` function.  When
specified, it is used to get the key for each incoming message, and
the latest value for each key is cached and can be retrieved
separately.

Signed-off-by: Sahas Subramanian <[email protected]>

Author: shsms

src/frequenz/channels/_latest_value_cache.py

@@ -8,12 +8,20 @@
 
 [LatestValueCache][frequenz.channels.LatestValueCache] takes a
 [Receiver][frequenz.channels.Receiver] as an argument and stores the latest
-value received by that receiver.  As soon as a value is received, its
+value received by that receiver.  It also takes an optional `key` function
+that allows you to group the values by a specific key.  If the `key` is
+provided, the cache will store the latest value for each key separately,
+otherwise it will store only the latest value received overall.
+
+As soon as a value is received, its
 [`has_value`][frequenz.channels.LatestValueCache.has_value] method returns
 `True`, and its [`get`][frequenz.channels.LatestValueCache.get] method returns
 the latest value received.  The `get` method will raise an exception if called
 before any messages have been received from the receiver.
 
+Both `has_value` and `get` methods can take an optional `key` argument to
+check or retrieve the latest value for that specific key.
+
 Example:
 ```python
 from frequenz.channels import Broadcast, LatestValueCache
@@ -32,12 +40,15 @@
 ```
 """
 
+from __future__ import annotations
+
 import asyncio
 import typing
 
 from ._receiver import Receiver
 
 T_co = typing.TypeVar("T_co", covariant=True)
+HashableT = typing.TypeVar("HashableT", bound=typing.Hashable)
 
 
 class _Sentinel:
@@ -48,15 +59,59 @@ def __str__(self) -> str:
         return "<no value received yet>"
 
 
-class LatestValueCache(typing.Generic[T_co]):
+_SENTINEL = _Sentinel()
+
+
+class LatestValueCache(typing.Generic[T_co, HashableT]):
     """A cache that stores the latest value in a receiver.
 
     It provides a way to look up the latest value in a stream without any delay,
     as long as there has been one value received.
     """
 
+    @typing.overload
+    def __init__(
+        self: LatestValueCache[T_co, _Sentinel],
+        receiver: Receiver[T_co],
+        *,
+        unique_id: str | None = None,
+        key: _Sentinel = _SENTINEL,
+    ) -> None:
+        """Create a new cache that does not use keys.
+
+        Args:
+            receiver: The receiver to cache.
+            unique_id: A string to help uniquely identify this instance. If not
+                provided, a unique identifier will be generated from the object's
+                [`id()`][id]. It is used mostly for debugging purposes.
+            key: This parameter is ignored when set to `None`.
+        """
+
+    @typing.overload
     def __init__(
-        self, receiver: Receiver[T_co], *, unique_id: str | None = None
+        self: LatestValueCache[T_co, HashableT],
+        receiver: Receiver[T_co],
+        *,
+        unique_id: str | None = None,
+        key: typing.Callable[[T_co], HashableT],
+    ) -> None:
+        """Create a new cache that uses keys.
+
+        Args:
+            receiver: The receiver to cache.
+            unique_id: A string to help uniquely identify this instance. If not
+                provided, a unique identifier will be generated from the object's
+                [`id()`][id]. It is used mostly for debugging purposes.
+            key: A function that takes a value and returns a key to group the values by.
+                If provided, the cache will store the latest value for each key separately.
+        """
+
+    def __init__(
+        self,
+        receiver: Receiver[T_co],
+        *,
+        unique_id: str | None = None,
+        key: typing.Callable[[T_co], typing.Any] | _Sentinel = _SENTINEL,
     ) -> None:
         """Create a new cache.
 
@@ -65,10 +120,16 @@ def __init__(
             unique_id: A string to help uniquely identify this instance. If not
                 provided, a unique identifier will be generated from the object's
                 [`id()`][id]. It is used mostly for debugging purposes.
+            key: An optional function that takes a value and returns a key to group the
+                values by. If provided, the cache will store the latest value for each
+                key separately. If not provided, it will store only the latest value
+                received overall.
         """
         self._receiver = receiver
+        self._key: typing.Callable[[T_co], HashableT] | _Sentinel = key
         self._unique_id: str = hex(id(self)) if unique_id is None else unique_id
-        self._latest_value: T_co | _Sentinel = _Sentinel()
+        self._latest_value: T_co | _Sentinel = _SENTINEL
+        self._latest_value_by_key: dict[HashableT, T_co] = {}
         self._task = asyncio.create_task(
             self._run(), name=f"LatestValueCache«{self._unique_id}»"
         )
@@ -78,34 +139,53 @@ def unique_id(self) -> str:
         """The unique identifier of this instance."""
         return self._unique_id
 
-    def get(self) -> T_co:
+    def get(self, key: HashableT | _Sentinel = _SENTINEL) -> T_co:
         """Return the latest value that has been received.
 
         This raises a `ValueError` if no value has been received yet. Use `has_value` to
         check whether a value has been received yet, before trying to access the value,
         to avoid the exception.
 
+        Args:
+            key: An optional key to retrieve the latest value for that key. If not
+                provided, it retrieves the latest value received overall.
+
         Returns:
             The latest value that has been received.
 
         Raises:
             ValueError: If no value has been received yet.
         """
+        if not isinstance(key, _Sentinel):
+            if key not in self._latest_value_by_key:
+                raise ValueError(f"No value received for key: {key!r}")
+            return self._latest_value_by_key[key]
+
         if isinstance(self._latest_value, _Sentinel):
             raise ValueError("No value has been received yet.")
         return self._latest_value
 
-    def has_value(self) -> bool:
+    def has_value(self, key: HashableT | _Sentinel = _SENTINEL) -> bool:
         """Check whether a value has been received yet.
 
+        If `key` is provided, it checks whether a value has been received for that key.
+
+        Args:
+            key: An optional key to check if a value has been received for that key.
+
         Returns:
             `True` if a value has been received, `False` otherwise.
         """
+        if not isinstance(key, _Sentinel):
+            return key in self._latest_value_by_key
         return not isinstance(self._latest_value, _Sentinel)
 
     async def _run(self) -> None:
         async for value in self._receiver:
             self._latest_value = value
+            if not isinstance(self._key, _Sentinel):
+                key = self._key(value)
+                self._latest_value_by_key[key] = value
 
     async def stop(self) -> None:
         """Stop the cache."""