53 lines
2.0 KiB
Python
53 lines
2.0 KiB
Python
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
|