1from typing import Any
2
3from errbot.storage.base import StorageBase, StoragePluginBase
4
5ROOTS = {}  # make a little bit of an emulated persistence.
6
7
8class MemoryStorage(StorageBase):
9    def __init__(self, namespace):
10        self.namespace = namespace
11        self.root = ROOTS.get(namespace, {})
12
13    def get(self, key: str) -> Any:
14        if key not in self.root:
15            raise KeyError(f"{key} doesn't exist.")
16        return self.root[key]
17
18    def set(self, key: str, value: Any) -> None:
19        self.root[key] = value
20
21    def remove(self, key: str):
22        if key not in self.root:
23            raise KeyError(f"{key} doesn't exist.")
24        del self.root[key]
25
26    def len(self):
27        return len(self.root)
28
29    def keys(self):
30        return self.root.keys()
31
32    def close(self) -> None:
33        ROOTS[self.namespace] = self.root
34
35
36class MemoryStoragePlugin(StoragePluginBase):
37    def open(self, namespace: str) -> StorageBase:
38        return MemoryStorage(namespace)
39