Add setdefault() method to Store

Can be useful in some cases.
This commit is contained in:
Ran Benita 2020-03-05 23:24:01 +02:00
parent 678d65f051
commit b1d7a187f2
2 changed files with 17 additions and 0 deletions

View File

@ -104,6 +104,15 @@ class Store:
except KeyError:
return default
def setdefault(self, key: StoreKey[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: StoreKey[T]) -> None:
"""Delete the value for key.

View File

@ -37,6 +37,14 @@ def test_store() -> None:
with pytest.raises(KeyError):
store[key1]
# setdefault
store[key1] = "existing"
assert store.setdefault(key1, "default") == "existing"
assert store[key1] == "existing"
key_setdefault = StoreKey[bytes]()
assert store.setdefault(key_setdefault, b"default") == b"default"
assert store[key_setdefault] == b"default"
# Can't accidentally add attributes to store object itself.
with pytest.raises(AttributeError):
store.foo = "nope" # type: ignore[attr-defined] # noqa: F821