pytest2/src/_pytest/stash.py

130 lines
3.7 KiB
Python

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 ``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``:
.. code-block:: python
stash: Stash = some_object.stash
If a module wants to store data in this Stash, it creates ``StashKey``\s
for its keys (at the module level):
.. code-block:: python
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]
Why use this?
-------------
Problem: module Internal defines an object. Module External, which
module Internal doesn't know about, receives the object and wants to
attach information to it, to be retrieved later given the object.
Bad solution 1: Module External assigns private attributes directly on
the object. This doesn't work well because the type checker doesn't
know about these attributes and it complains about undefined attributes.
Bad solution 2: module Internal adds a ``Dict[str, Any]`` attribute to
the object. Module External stores its data in private keys of this dict.
This doesn't work well because retrieved values are untyped.
Good solution: module Internal adds a ``Stash`` to the object. Module
External mints StashKeys for its own keys. Module External stores and
retrieves its data using these keys.
"""
__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)