stash.py 3.0 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116
  1. from __future__ import annotations
  2. from typing import Any
  3. from typing import cast
  4. from typing import Generic
  5. from typing import TypeVar
  6. __all__ = ["Stash", "StashKey"]
  7. T = TypeVar("T")
  8. D = TypeVar("D")
  9. class StashKey(Generic[T]):
  10. """``StashKey`` is an object used as a key to a :class:`Stash`.
  11. A ``StashKey`` is associated with the type ``T`` of the value of the key.
  12. A ``StashKey`` is unique and cannot conflict with another key.
  13. .. versionadded:: 7.0
  14. """
  15. __slots__ = ()
  16. class Stash:
  17. r"""``Stash`` is a type-safe heterogeneous mutable mapping that
  18. allows keys and value types to be defined separately from
  19. where it (the ``Stash``) is created.
  20. Usually you will be given an object which has a ``Stash``, for example
  21. :class:`~pytest.Config` or a :class:`~_pytest.nodes.Node`:
  22. .. code-block:: python
  23. stash: Stash = some_object.stash
  24. If a module or plugin wants to store data in this ``Stash``, it creates
  25. :class:`StashKey`\s for its keys (at the module level):
  26. .. code-block:: python
  27. # At the top-level of the module
  28. some_str_key = StashKey[str]()
  29. some_bool_key = StashKey[bool]()
  30. To store information:
  31. .. code-block:: python
  32. # Value type must match the key.
  33. stash[some_str_key] = "value"
  34. stash[some_bool_key] = True
  35. To retrieve the information:
  36. .. code-block:: python
  37. # The static type of some_str is str.
  38. some_str = stash[some_str_key]
  39. # The static type of some_bool is bool.
  40. some_bool = stash[some_bool_key]
  41. .. versionadded:: 7.0
  42. """
  43. __slots__ = ("_storage",)
  44. def __init__(self) -> None:
  45. self._storage: dict[StashKey[Any], object] = {}
  46. def __setitem__(self, key: StashKey[T], value: T) -> None:
  47. """Set a value for key."""
  48. self._storage[key] = value
  49. def __getitem__(self, key: StashKey[T]) -> T:
  50. """Get the value for key.
  51. Raises ``KeyError`` if the key wasn't set before.
  52. """
  53. return cast(T, self._storage[key])
  54. def get(self, key: StashKey[T], default: D) -> T | D:
  55. """Get the value for key, or return default if the key wasn't set
  56. before."""
  57. try:
  58. return self[key]
  59. except KeyError:
  60. return default
  61. def setdefault(self, key: StashKey[T], default: T) -> T:
  62. """Return the value of key if already set, otherwise set the value
  63. of key to default and return default."""
  64. try:
  65. return self[key]
  66. except KeyError:
  67. self[key] = default
  68. return default
  69. def __delitem__(self, key: StashKey[T]) -> None:
  70. """Delete the value for key.
  71. Raises ``KeyError`` if the key wasn't set before.
  72. """
  73. del self._storage[key]
  74. def __contains__(self, key: StashKey[T]) -> bool:
  75. """Return whether key was set."""
  76. return key in self._storage
  77. def __len__(self) -> int:
  78. """Return how many items exist in the stash."""
  79. return len(self._storage)