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

8

9__all__ = ["Stash", "StashKey"]

10

11

12 T = TypeVar("T")

13 D = TypeVar("D")

14

15

16 class StashKey(Generic[T]):

17 """``StashKey`` is an object used as a key to a :class:`Stash`.

18

19 A ``StashKey`` is associated with the type ``T`` of the value of the key.

20

21 A ``StashKey`` is unique and cannot conflict with another key.

22 """

23

24 __slots__ = ()

25

26

27 class Stash:

28 r"""``Stash`` is a type-safe heterogeneous mutable mapping that

29 allows keys and value types to be defined separately from

30 where it (the ``Stash``) is created.

31

32 Usually you will be given an object which has a ``Stash``, for example

33 :class:`~pytest.Config` or a :class:`~_pytest.nodes.Node`:

34

35 .. code-block:: python

36

37 stash: Stash = some_object.stash

38

39 If a module or plugin wants to store data in this ``Stash``, it creates

40 :class:`StashKey`\s for its keys (at the module level):

41

42 .. code-block:: python

43

44 # At the top-level of the module

45 some_str_key = StashKey[str]()

46 some_bool_key = StashKey[bool]()

47

48 To store information:

49

50 .. code-block:: python

51

52 # Value type must match the key.

53 stash[some_str_key] = "value"

54 stash[some_bool_key] = True

55

56 To retrieve the information:

57

58 .. code-block:: python

59

60 # The static type of some_str is str.

61 some_str = stash[some_str_key]

62 # The static type of some_bool is bool.

63 some_bool = stash[some_bool_key]

64 """

65

66 __slots__ = ("_storage",)

67

68 def __init__(self) -> None:

69 self._storage: Dict[StashKey[Any], object] = {}

70

71 def __setitem__(self, key: StashKey[T], value: T) -> None:

72 """Set a value for key."""

73 self._storage[key] = value

74

75 def __getitem__(self, key: StashKey[T]) -> T:

76 """Get the value for key.

77

78 Raises ``KeyError`` if the key wasn't set before.

79 """

80 return cast(T, self._storage[key])

81

82 def get(self, key: StashKey[T], default: D) -> Union[T, D]:

83 """Get the value for key, or return default if the key wasn't set

84 before."""

85 try:

86 return self[key]

87 except KeyError:

88 return default

89

90 def setdefault(self, key: StashKey[T], default: T) -> T:

91 """Return the value of key if already set, otherwise set the value

92 of key to default and return default."""

93 try:

94 return self[key]

95 except KeyError:

96 self[key] = default

97 return default

98

99 def __delitem__(self, key: StashKey[T]) -> None:

100 """Delete the value for key.

101

102 Raises ``KeyError`` if the key wasn't set before.

103 """

104 del self._storage[key]

105

106 def __contains__(self, key: StashKey[T]) -> bool:

107 """Return whether key was set."""

108 return key in self._storage

109

110 def __len__(self) -> int:

111 """Return how many items exist in the stash."""

112 return len(self._storage)