diff --git a/CHANGES.rst b/CHANGES.rst index 320e4836c..5d480438e 100644 --- a/CHANGES.rst +++ b/CHANGES.rst @@ -9,6 +9,7 @@ Unreleased - ``CacheControl.no_transform`` is a boolean when present. ``min_fresh`` is ``None`` when not present. Added the ``must_understand`` attribute. Fixed some typing issues on cache control. :issue:`2881` +- Added ``TypeConversionDict.pop`` method. :issue:`2883` Version 3.0.6 diff --git a/src/werkzeug/datastructures/structures.py b/src/werkzeug/datastructures/structures.py index 4279ceb98..90479091c 100644 --- a/src/werkzeug/datastructures/structures.py +++ b/src/werkzeug/datastructures/structures.py @@ -88,6 +88,54 @@ def get(self, key, default=None, type=None): rv = default return rv + def pop(self, key, default=_missing, type=None): + """Like :meth:`get` but removes the key/value pair. + + >>> d = TypeConversionDict(foo='42', bar='blub') + >>> d.pop('foo', type=int) + 42 + >>> 'foo' in d + False + >>> d.pop('bar', -1, type=int) + -1 + >>> 'bar' in d + False + + :param key: The key to be looked up. + :param default: The default value to be returned if the key is not + in the dictionary. If not further specified it's + an :exc:`KeyError`. + :param type: A callable that is used to cast the value in the dict. + If a :exc:`ValueError` or a :exc:`TypeError` is raised + by this callable the default value is returned. + + .. admonition:: note + + If the type conversion fails, the key is **not** removed from the + dictionary. + + """ + try: + rv = self[key] + except KeyError: + if default is _missing: + raise + return default + if type is not None: + try: + rv = type(rv) + except (ValueError, TypeError): + if default is _missing: + return None + return default + try: + # This method is not meant to be thread-safe, but at least lets not + # fall over if the dict was mutated between the get and the delete. -MK + del self[key] + except KeyError: + pass + return rv + class ImmutableTypeConversionDict(ImmutableDictMixin, TypeConversionDict): """Works like a :class:`TypeConversionDict` but does not support diff --git a/src/werkzeug/datastructures/structures.pyi b/src/werkzeug/datastructures/structures.pyi index 7086ddae1..ef055f4af 100644 --- a/src/werkzeug/datastructures/structures.pyi +++ b/src/werkzeug/datastructures/structures.pyi @@ -36,10 +36,28 @@ class TypeConversionDict(dict[K, V]): def get(self, key: K, default: D, type: Callable[[V], T]) -> D | T: ... @overload def get(self, key: K, type: Callable[[V], T]) -> T | None: ... + @overload + def pop(self, key: K) -> V: ... + @overload + def pop(self, key: K, default: V) -> V: ... + @overload + def pop(self, key: K, default: D) -> V | D: ... + @overload + def pop(self, key: K, default: D, type: Callable[[V], T]) -> D | T: ... + @overload + def pop(self, key: K, type: Callable[[V], T]) -> T | None: ... class ImmutableTypeConversionDict(ImmutableDictMixin[K, V], TypeConversionDict[K, V]): def copy(self) -> TypeConversionDict[K, V]: ... def __copy__(self) -> ImmutableTypeConversionDict[K, V]: ... + @overload + def pop(self, key: K, default: V | None = ...) -> NoReturn: ... + @overload + def pop(self, key: K, default: D) -> NoReturn: ... + @overload + def pop(self, key: K, default: D, type: Callable[[V], T]) -> NoReturn: ... + @overload + def pop(self, key: K, type: Callable[[V], T]) -> NoReturn: ... class MultiDict(TypeConversionDict[K, V]): def __init__( @@ -74,6 +92,14 @@ class MultiDict(TypeConversionDict[K, V]): @overload def pop(self, key: K) -> V: ... @overload + def pop(self, key: K, default: V) -> V: ... + @overload + def pop(self, key: K, default: D) -> V | D: ... + @overload + def pop(self, key: K, default: D, type: Callable[[V], T]) -> D | T: ... + @overload + def pop(self, key: K, type: Callable[[V], T]) -> T | None: ... + @overload def pop(self, key: K, default: V | T = ...) -> V | T: ... def popitem(self) -> tuple[K, V]: ... def poplist(self, key: K) -> list[V]: ... @@ -119,6 +145,14 @@ class OrderedMultiDict(MultiDict[K, V]): @overload def pop(self, key: K) -> V: ... @overload + def pop(self, key: K, default: V) -> V: ... + @overload + def pop(self, key: K, default: D) -> V | D: ... + @overload + def pop(self, key: K, default: D, type: Callable[[V], T]) -> D | T: ... + @overload + def pop(self, key: K, type: Callable[[V], T]) -> T | None: ... + @overload def pop(self, key: K, default: V | T = ...) -> V | T: ... def popitem(self) -> tuple[K, V]: ... def popitemlist(self) -> tuple[K, list[V]]: ... diff --git a/tests/test_datastructures.py b/tests/test_datastructures.py index 830dfefd5..24247e738 100644 --- a/tests/test_datastructures.py +++ b/tests/test_datastructures.py @@ -304,6 +304,19 @@ def test_dict_is_hashable(self): assert immutable in x assert immutable2 in x + def test_get_does_not_raise(self): + cls = self.storage_class + immutable = cls({"a": 1}) + assert immutable.get("a") == 1 + + def test_pop_raises(self): + cls = self.storage_class + immutable = cls({"a": 1}) + with pytest.raises(TypeError): + immutable.pop("a") + with pytest.raises(TypeError): + immutable.popitem() + class TestImmutableTypeConversionDict(_ImmutableDictTests): storage_class = ds.ImmutableTypeConversionDict @@ -545,20 +558,56 @@ def test_get_description(self): class TestTypeConversionDict: storage_class = ds.TypeConversionDict - def test_value_conversion(self): + class MyException(Exception): + def raize(self, *args, **kwargs): + raise self + + def test_get_value_conversion(self): d = self.storage_class(foo="1") assert d.get("foo", type=int) == 1 - def test_return_default_when_conversion_is_not_possible(self): + def test_get_return_default_when_conversion_is_not_possible(self): d = self.storage_class(foo="bar", baz=None) assert d.get("foo", default=-1, type=int) == -1 assert d.get("baz", default=-1, type=int) == -1 - def test_propagate_exceptions_in_conversion(self): + def test_get_propagate_exceptions_in_conversion(self): + d = self.storage_class(foo="bar") + with pytest.raises(self.MyException): + d.get("foo", type=lambda x: self.MyException().raize()) + + def test_get_error_in_conversion(self): d = self.storage_class(foo="bar") - switch = {"a": 1} + assert d.get("foo", type=int) is None + + def test_pop_value_conversion(self): + d = self.storage_class(foo="1") + assert d.pop("foo", type=int) == 1 + assert "foo" not in d + + def test_pop_return_when_conversion_is_not_possible(self): + d = self.storage_class(foo="bar", baz=None) + assert d.pop("foo", type=int) is None + assert "foo" in d # key is still in the dict, because the conversion failed + assert d.pop("baz", type=int) is None + assert "baz" in d # key is still in the dict, because the conversion failed + + def test_pop_return_default_when_conversion_is_not_possible(self): + d = self.storage_class(foo="bar", baz=None) + assert d.pop("foo", default=-1, type=int) == -1 + assert "foo" in d # key is still in the dict, because the conversion failed + assert d.pop("baz", default=-1, type=int) == -1 + assert "baz" in d # key is still in the dict, because the conversion failed + + def test_pop_propagate_exceptions_in_conversion(self): + d = self.storage_class(foo="bar") + with pytest.raises(self.MyException): + d.pop("foo", type=lambda x: self.MyException().raize()) + + def test_pop_key_error(self): + d = self.storage_class() with pytest.raises(KeyError): - d.get("foo", type=lambda x: switch[x]) + d.pop("foo") class TestCombinedMultiDict: