aboutsummaryrefslogtreecommitdiffstats
path: root/contrib/python/pytest/py3/_pytest/stash.py
blob: e61d75b95f7d93120abbf3266eb52cd1b78aae8e (plain) (blame)
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
from typing import Any
from typing import cast
from typing import Dict
from typing import Generic
from typing import TypeVar
from typing import Union


__all__ = ["Stash", "StashKey"]


T = TypeVar("T")
D = TypeVar("D")


class StashKey(Generic[T]):
    """``StashKey`` is an object used as a key to a :class:`Stash`.

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

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

    __slots__ = ()


class Stash:
    r"""``Stash`` is a type-safe heterogeneous mutable mapping that
    allows keys and value types to be defined separately from
    where it (the ``Stash``) is created.

    Usually you will be given an object which has a ``Stash``, for example
    :class:`~pytest.Config` or a :class:`~_pytest.nodes.Node`:

    .. code-block:: python

        stash: Stash = some_object.stash

    If a module or plugin wants to store data in this ``Stash``, it creates
    :class:`StashKey`\s for its keys (at the module level):

    .. code-block:: python

        # At the top-level of the module
        some_str_key = StashKey[str]()
        some_bool_key = StashKey[bool]()

    To store information:

    .. code-block:: python

        # Value type must match the key.
        stash[some_str_key] = "value"
        stash[some_bool_key] = True

    To retrieve the information:

    .. code-block:: python

        # The static type of some_str is str.
        some_str = stash[some_str_key]
        # The static type of some_bool is bool.
        some_bool = stash[some_bool_key]
    """

    __slots__ = ("_storage",)

    def __init__(self) -> None:
        self._storage: Dict[StashKey[Any], object] = {}

    def __setitem__(self, key: StashKey[T], value: T) -> None:
        """Set a value for key."""
        self._storage[key] = value

    def __getitem__(self, key: StashKey[T]) -> T:
        """Get the value for key.

        Raises ``KeyError`` if the key wasn't set before.
        """
        return cast(T, self._storage[key])

    def get(self, key: StashKey[T], default: D) -> Union[T, D]:
        """Get the value for key, or return default if the key wasn't set
        before."""
        try:
            return self[key]
        except KeyError:
            return default

    def setdefault(self, key: StashKey[T], default: T) -> T:
        """Return the value of key if already set, otherwise set the value
        of key to default and return default."""
        try:
            return self[key]
        except KeyError:
            self[key] = default
            return default

    def __delitem__(self, key: StashKey[T]) -> None:
        """Delete the value for key.

        Raises ``KeyError`` if the key wasn't set before.
        """
        del self._storage[key]

    def __contains__(self, key: StashKey[T]) -> bool:
        """Return whether key was set."""
        return key in self._storage

    def __len__(self) -> int:
        """Return how many items exist in the stash."""
        return len(self._storage)