import six
from .generic import KeyStorage, UnknownKeyName
if six.PY2:
    from collections import Mapping
else:
    from collections.abc import Mapping


__all__ = [
    'MemoryKeyStorage',
]


class MemoryKeyStorage(KeyStorage):
    """ In-memory keys storage """

    def __init__(self, keys):
        super(MemoryKeyStorage, self).__init__()
        self._keys = {}
        for key, value in six.iteritems(keys):
            if isinstance(value, Mapping) and not isinstance(value, KeyStorage):
                self._keys[key] = MemoryKeyStorage(keys=value)
            else:
                self._keys[key] = value

    def __getitem__(self, key):
        try:
            return self._keys.__getitem__(key)
        except KeyError:
            raise UnknownKeyName(key)

    def __iter__(self, *args, **kwargs):
        return self._keys.__iter__(*args, **kwargs)

    def __len__(self, *args, **kwargs):
        return self._keys.__len__(*args, **kwargs)
