diff --git a/fsspec/mapping.py b/fsspec/mapping.py index 26bd22099..41386ff30 100644 --- a/fsspec/mapping.py +++ b/fsspec/mapping.py @@ -9,6 +9,7 @@ from fsspec.core import url_to_fs logger = logging.getLogger("fsspec.mapping") +_MISSING = object() class FSMap(MutableMapping): @@ -161,9 +162,14 @@ def __getitem__(self, key, default=None): raise KeyError(key) from exc return result - def pop(self, key, default=None): + def pop(self, key, default=_MISSING): """Pop data""" - result = self.__getitem__(key, default) + try: + result = self[key] + except KeyError: + if default is _MISSING: + raise + return default try: del self[key] except KeyError: diff --git a/fsspec/tests/test_mapping.py b/fsspec/tests/test_mapping.py index 075b6cc23..ea7f6cf83 100644 --- a/fsspec/tests/test_mapping.py +++ b/fsspec/tests/test_mapping.py @@ -239,3 +239,34 @@ def test_fsmap_dirfs(): fs = m.dirfs assert isinstance(fs, fsspec.implementations.dirfs.DirFileSystem) assert fs.path == m.root + + +@pytest.mark.parametrize("protocol", ["memory", "file"]) +@pytest.mark.parametrize("default", [None, False, 0, b"", []]) +def test_pop_missing_key_returns_explicit_default(tmp_path, protocol, default): + mapper = fsspec.get_mapper( + f"{protocol}://{tmp_path.as_posix()}/mapping", create=True + ) + + assert mapper.pop("missing", default) is default + + +@pytest.mark.parametrize("protocol", ["memory", "file"]) +def test_pop_without_default_raises_for_missing_key(tmp_path, protocol): + mapper = fsspec.get_mapper( + f"{protocol}://{tmp_path.as_posix()}/mapping", create=True + ) + + with pytest.raises(KeyError, match="missing"): + mapper.pop("missing") + + +@pytest.mark.parametrize("protocol", ["memory", "file"]) +def test_pop_existing_key_returns_and_removes_value(tmp_path, protocol): + mapper = fsspec.get_mapper( + f"{protocol}://{tmp_path.as_posix()}/mapping", create=True + ) + mapper["key"] = b"value" + + assert mapper.pop("key", None) == b"value" + assert "key" not in mapper