from functools import reduce
from operator import getitem
from typing import Any, Callable, Sequence, Tuple
from basic_utils.seq_helpers import butlast, last
__all__ = [
'get_in_dict',
'get_keys',
'prune_dict',
'set_in_dict',
]
[docs]def get_keys(d: dict, keys: Sequence[Any], default: Callable=None) -> Tuple:
"""
Returns multiple values for keys in a dictionary
Empty key values will be None by default
>>> d = {'x': 24, 'y': 25}
>>> get_keys(d, ('x', 'y', 'z'))
(24, 25, None)
"""
return tuple(d.get(key, default) for key in keys)
[docs]def get_in_dict(d: dict, keys: Sequence[str]) -> Any:
"""
Retrieve nested key from dictionary
>>> d = {'a': {'b': {'c': 3}}}
>>> get_in_dict(d, ('a', 'b', 'c'))
3
"""
return reduce(getitem, keys, d)
[docs]def set_in_dict(d: dict, keys: Sequence[str], value: Any) -> None:
"""
Sets a value inside a nested dictionary
>>> d = {'a': {'b': {'c': 3}}}
>>> set_in_dict(d, ('a', 'b', 'c'), 10)
>>> d
{'a': {'b': {'c': 10}}}
"""
get_in_dict(d, butlast(keys))[last(keys)] = value
[docs]def prune_dict(d: dict) -> dict:
"""
Returns new dictionary with falesly values removed.
>>> prune_dict({'a': [], 'b': 2, 'c': False})
{'b': 2}
"""
return filter_values(d, bool)
def filter_keys(d: dict, predicate: Callable) -> dict:
"""
Returns new subset dict of d where keys pass predicate fn
>>> filter_keys({'Lisa': 8, 'Marge': 36}, lambda x: len(x) > 4)
{'Marge': 36}
"""
return {k: v for k, v in d.items() if predicate(k)}
def filter_values(d: dict, predicate: Callable) -> dict:
"""
Returns new subset dict of d where values pass predicate fn
>>> d = {'Homer': 39, 'Marge': 36, 'Bart': 10}
>>> filter_values(d, lambda x: x < 20)
{'Bart': 10}
"""
return {k: v for k, v in d.items() if predicate(v)}