From b1d7a187f29de08152c174e1468c8b5743c96dd3 Mon Sep 17 00:00:00 2001
From: Ran Benita <ran@unusedvar.com>
Date: Thu, 5 Mar 2020 23:24:01 +0200
Subject: [PATCH] Add setdefault() method to Store

Can be useful in some cases.
---
 src/_pytest/store.py  | 9 +++++++++
 testing/test_store.py | 8 ++++++++
 2 files changed, 17 insertions(+)

diff --git a/src/_pytest/store.py b/src/_pytest/store.py
index 0dcea1b93..eed50d103 100644
--- a/src/_pytest/store.py
+++ b/src/_pytest/store.py
@@ -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.
 
diff --git a/testing/test_store.py b/testing/test_store.py
index ae6cfccdd..98014887e 100644
--- a/testing/test_store.py
+++ b/testing/test_store.py
@@ -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