diff --git a/benedict/dicts/base/base_dict.py b/benedict/dicts/base/base_dict.py index 887191f9..8452768f 100644 --- a/benedict/dicts/base/base_dict.py +++ b/benedict/dicts/base/base_dict.py @@ -9,7 +9,7 @@ Reversible, ValuesView, ) -from typing import Any, Self +from typing import Any, Self, overload from typing_extensions import TypeVar @@ -17,6 +17,7 @@ _K = TypeVar("_K", default=str) _V = TypeVar("_V", default=Any) +_D = TypeVar("_D") class BaseDict(dict[_K, _V]): @@ -207,7 +208,16 @@ def dict(self) -> dict[_K, _V]: return self._dict return self - def get(self, key: _K, default: _V | None = None) -> _V | None: # type: ignore[override] + @overload + def get(self, key: _K, default: None = None) -> _V | None: ... + + @overload + def get(self, key: _K, default: _V) -> _V: ... + + @overload + def get(self, key: _K, default: _D) -> _V | _D: ... + + def get(self, key: _K, default: Any = None) -> Any: if self._dict is not None: return self._dict.get(key, default) return super().get(key, default) @@ -222,11 +232,20 @@ def keys(self) -> KeysView[_K]: # type: ignore[override] return self._dict.keys() return super().keys() - def pop(self, key: _K, *args: Any) -> _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: ... + + def pop(self, key: _K, *args: Any) -> Any: self._check_frozen() if self._dict is not None: - return self._dict.pop(key, *args) # type: ignore[no-any-return] - return super().pop(key, *args) # type: ignore[no-any-return] + return self._dict.pop(key, *args) + return super().pop(key, *args) def popitem(self) -> tuple[_K, _V]: self._check_frozen() diff --git a/tests/dicts/base/test_base_dict.py b/tests/dicts/base/test_base_dict.py index cc955a31..56f1aaf2 100644 --- a/tests/dicts/base/test_base_dict.py +++ b/tests/dicts/base/test_base_dict.py @@ -39,6 +39,23 @@ def test_init_readonly_mapping_with_keywords(self) -> None: self.assertEqual(source, {"existing": 1}) self.assertIs(BaseDict(source).dict(), source) + def test_get_and_pop_preserve_default_types(self) -> None: + values = [BaseDict[str, int](found=3), BaseDict[str, int]({"found": 3})] + for value in values: + with self.subTest(wrapped=value._dict is not None): + assert_type(value.get("missing", "fallback"), int | str) + assert_type(value.get("missing"), int | None) + assert_type(value.get("found", 0), int) + self.assertEqual(value.get("missing", "fallback"), "fallback") + self.assertEqual(value.get("found", "fallback"), 3) + assert_type(value.pop("missing", "fallback"), int | str) + assert_type(value.pop("missing", None), int | None) + found = value.pop("found") + assert_type(found, int) + self.assertEqual(value.pop("missing", "fallback"), "fallback") + self.assertIsNone(value.pop("missing", None)) + self.assertEqual(found, 3) + def test_popitem(self) -> None: b = BaseDict(a=1, b=2) self.assertEqual(b.popitem(), ("b", 2))