Spaces:
Running on Zero
Running on Zero
File size: 1,478 Bytes
0122a25 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 | """This module contains dictionary utility functions."""
from __future__ import annotations
from typing import Any, NamedTuple
def get_all_keys(entry: NamedTuple) -> list[str]:
"""Get all keys in a NamedTuple."""
keys = []
for key in entry._fields:
if is_namedtuple(getattr(entry, key)):
keys.extend(
[f"{key}.{k}" for k in get_all_keys(getattr(entry, key))]
)
else:
keys.append(key)
return keys
def get_from_namedtuple(entry: NamedTuple, key: str) -> Any: # type: ignore
"""Get a value from a nested Named tuple.
Example passing key = "test.my.data" will resolve the value of the
named tuple at 'test' 'my' 'data'.
Raises:
ValueError: If the key is not present in the named tuple.
"""
keys = key.split(".")
first_key = keys[0]
if not hasattr(entry, first_key):
raise ValueError(
f"Key {first_key} not in named tuple! Current keys: "
f"{get_all_keys(entry)}"
)
if len(keys) == 1:
return getattr(entry, first_key)
return get_from_namedtuple(getattr(entry, first_key), ".".join(keys[1:]))
def is_namedtuple(obj: object) -> bool:
"""Check if obj is namedtuple.
https://github.com/pytorch/pytorch/blob/v1.8.1/torch/nn/parallel/scatter_gather.py#L4-L8
"""
return (
isinstance(obj, tuple)
and hasattr(obj, "_asdict")
and hasattr(obj, "_fields")
)
|