map_of_maps_attribute.py

from typing import Any
from typing import Dict
from typing import Type
from typing import TYPE_CHECKING
from typing import TypeVar
​
from pynamodb.attributes import MapAttribute
​
​
MapT = TypeVar('MapT', bound=MapAttribute)
​
​
class MapOfMapsAttribute(MapAttribute[str, MapT]):
    """
    Like a raw MapAttribute, but you must provide an 'of' keyword argument to the constructor,
    and all attributes are assumed to be maps.
    """
    _element_type: Type[MapAttribute]
​
    # Ensures that 'of' kwarg is persisted when used as a non-container.
    attribute_args = [*MapAttribute.attribute_args, 'of']
​
    def is_raw(self) -> bool:  # type:ignore[override]
        return True
​
    def deserialize(self, values: Dict[str, Any]) -> Dict[str, Any]:
        element_attr = self._element_type()
        return {
            k: element_attr.deserialize(attr_value)
            for k, v in values.items()
            for attr_type, attr_value in v.items()
            if attr_type == element_attr.attr_type
        }
​
    def _make_attribute(self) -> None:
        try:
            element_type = self.attribute_kwargs.pop('of')
        except KeyError:
            raise ValueError(f"{type(self).__name__} must be defined with 'of' keyword argument'")
        super()._make_attribute()
        self._element_type = element_type
​
    if TYPE_CHECKING:
        def __init__(self, of: Type[MapT], **kwargs: Any) -> None:
            ...