from collections import defaultdict


class CustomStorage(object):
    __slots__ = ['_storage', 'full_storage']

    def __init__(self):
        self._storage = defaultdict(dict)
        self.full_storage = defaultdict(dict)

    def clear_data(self):
        self.__init__()

    def write_data(self, context, key, value, write_to_xml=True):
        self.full_storage[context][key] = value
        if write_to_xml:
            self._storage[context][key] = value

    def read_data(self, context, key):
        dic = self.full_storage[context]
        if hasattr(key, "__iter__"):
            # "key" is actually a list of keys to retrieve at once
            return map(dic.get, key)
        # single key to retrieve
        return dic.get(key)

    def delete_data(self, context, key):
        try:
            del self._storage[context][key]
            del self.full_storage[context][key]
        except KeyError:
            pass

    def locate_data(self, key):
        result = []
        for context in self._storage:
            if key in self._storage[context]:
                result.append(context)

        return result
