chore: import upstream snapshot with attribution

rllib/utils/lambda_defaultdict.py

@@ -0,0 +1,52 @@
+from collections import defaultdict
+from typing import Any, Callable
+
+
+class LambdaDefaultDict(defaultdict):
+    """A defaultdict that creates default values based on the associated key.
+
+    Note that the standard defaultdict can only produce default values (via its factory)
+    that are independent of the key under which they are stored.
+    As opposed to that, the lambda functions used as factories for this
+    `LambdaDefaultDict` class do accept a single argument: The missing key.
+    If a missing key is accessed by the user, the provided lambda function is called
+    with this missing key as its argument. The returned value is stored in the
+    dictionary under that key and returned.
+
+    Example:
+
+        In this example, if you try to access a key that doesn't exist, it will call
+        the lambda function, passing it the missing key. The function will return a
+        string, which will be stored in the dictionary under that key.
+
+        .. testcode::
+
+            from ray.rllib.utils.lambda_defaultdict import LambdaDefaultDict
+
+            default_dict = LambdaDefaultDict(lambda missing_key: f"Value for {missing_key}")
+            print(default_dict["a"])
+
+        .. testoutput::
+
+            Value for a
+    """  # noqa: E501
+
+    def __init__(self, default_factory: Callable[[str], Any], *args, **kwargs):
+        """Initializes a LambdaDefaultDict instance.
+
+        Args:
+            default_factory: The default factory callable, taking a string (key)
+                and returning the default value to use for that key.
+        """
+        if not callable(default_factory):
+            raise TypeError("First argument must be a Callable!")
+
+        # We will handle the factory in __missing__ method.
+        super().__init__(None, *args, **kwargs)
+
+        self.default_factory = default_factory
+
+    def __missing__(self, key):
+        # Call default factory with the key as argument.
+        self[key] = value = self.default_factory(key)
+        return value