1from typing import Any
2from typing import cast
3from typing import Dict
4from typing import Generic
5from typing import TypeVar
6from typing import Union
7
8
9__all__ = ["Store", "StoreKey"]
10
11
12T = TypeVar("T")
13D = TypeVar("D")
14
15
16class StoreKey(Generic[T]):
17    """StoreKey is an object used as a key to a Store.
18
19    A StoreKey is associated with the type T of the value of the key.
20
21    A StoreKey is unique and cannot conflict with another key.
22    """
23
24    __slots__ = ()
25
26
27class Store:
28    """Store is a type-safe heterogenous mutable mapping that
29    allows keys and value types to be defined separately from
30    where it (the Store) is created.
31
32    Usually you will be given an object which has a ``Store``:
33
34    .. code-block:: python
35
36        store: Store = some_object.store
37
38    If a module wants to store data in this Store, it creates StoreKeys
39    for its keys (at the module level):
40
41    .. code-block:: python
42
43        some_str_key = StoreKey[str]()
44        some_bool_key = StoreKey[bool]()
45
46    To store information:
47
48    .. code-block:: python
49
50        # Value type must match the key.
51        store[some_str_key] = "value"
52        store[some_bool_key] = True
53
54    To retrieve the information:
55
56    .. code-block:: python
57
58        # The static type of some_str is str.
59        some_str = store[some_str_key]
60        # The static type of some_bool is bool.
61        some_bool = store[some_bool_key]
62
63    Why use this?
64    -------------
65
66    Problem: module Internal defines an object. Module External, which
67    module Internal doesn't know about, receives the object and wants to
68    attach information to it, to be retrieved later given the object.
69
70    Bad solution 1: Module External assigns private attributes directly on
71    the object. This doesn't work well because the type checker doesn't
72    know about these attributes and it complains about undefined attributes.
73
74    Bad solution 2: module Internal adds a ``Dict[str, Any]`` attribute to
75    the object. Module External stores its data in private keys of this dict.
76    This doesn't work well because retrieved values are untyped.
77
78    Good solution: module Internal adds a ``Store`` to the object. Module
79    External mints StoreKeys for its own keys. Module External stores and
80    retrieves its data using these keys.
81    """
82
83    __slots__ = ("_store",)
84
85    def __init__(self) -> None:
86        self._store = {}  # type: Dict[StoreKey[Any], object]
87
88    def __setitem__(self, key: StoreKey[T], value: T) -> None:
89        """Set a value for key."""
90        self._store[key] = value
91
92    def __getitem__(self, key: StoreKey[T]) -> T:
93        """Get the value for key.
94
95        Raises ``KeyError`` if the key wasn't set before.
96        """
97        return cast(T, self._store[key])
98
99    def get(self, key: StoreKey[T], default: D) -> Union[T, D]:
100        """Get the value for key, or return default if the key wasn't set
101        before."""
102        try:
103            return self[key]
104        except KeyError:
105            return default
106
107    def setdefault(self, key: StoreKey[T], default: T) -> T:
108        """Return the value of key if already set, otherwise set the value
109        of key to default and return default."""
110        try:
111            return self[key]
112        except KeyError:
113            self[key] = default
114            return default
115
116    def __delitem__(self, key: StoreKey[T]) -> None:
117        """Delete the value for key.
118
119        Raises ``KeyError`` if the key wasn't set before.
120        """
121        del self._store[key]
122
123    def __contains__(self, key: StoreKey[T]) -> bool:
124        """Return whether key was set."""
125        return key in self._store
126