You can not select more than 25 topics Topics must start with a letter or number, can include dashes ('-') and can be up to 35 characters long.

126 lines
3.5KB

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