KeyMap (#110)

doc/source/changelog.rst

@@ -5,6 +5,7 @@ Changelog
 3.1.0 - Unreleased
 ------------------
 - Dropped support for Python 3.8 (:pr:`106`) `Guido Imperiale`_
+- New object :class:`KeyMap` (:pr:`110`) `Guido Imperiale`_
 
 
 3.0.0 - 2023-04-17

doc/source/index.rst

@@ -55,6 +55,8 @@ API
    :members:
 .. autoclass:: Func
    :members:
+.. autoclass:: KeyMap
+   :members:
 .. autoclass:: LMDB
    :members:
 .. autoclass:: LRU

zict/__init__.py

@@ -4,6 +4,7 @@
 from zict.cache import WeakValueMapping as WeakValueMapping
 from zict.file import File as File
 from zict.func import Func as Func
+from zict.keymap import KeyMap as KeyMap
 from zict.lmdb import LMDB as LMDB
 from zict.lru import LRU as LRU
 from zict.sieve import Sieve as Sieve

zict/func.py

@@ -9,7 +9,7 @@
 
 
 class Func(ZictBase[KT, VT], Generic[KT, VT, WT]):
-    """Buffer a MutableMapping with a pair of input/output functions
+    """Translate the values of a MutableMapping with a pair of input/output functions
 
     Parameters
     ----------
@@ -19,6 +19,10 @@ class Func(ZictBase[KT, VT], Generic[KT, VT, WT]):
         Function to call on value as we pull it from the mapping
     d: MutableMapping
 
+    See Also
+    --------
+    KeyMap
+
     Examples
     --------
     >>> def double(x):

zict/keymap.py

@@ -0,0 +1,81 @@
+from __future__ import annotations
+
+from collections.abc import Callable, Iterator, MutableMapping
+from typing import Generic, TypeVar
+
+from zict.common import KT, VT, ZictBase, close, discard, flush, locked
+
+JT = TypeVar("JT")
+
+
+class KeyMap(ZictBase[KT, VT], Generic[KT, JT, VT]):
+    """Translate the keys of a MutableMapping with a pair of input/output functions
+
+    Parameters
+    ----------
+    fn: callable
+        Function to call on a key of the KeyMap to transform it to a key of the wrapped
+        mapping. It must be pure (if called twice on the same key it must return
+        the same result) and it must not generate collisions. In other words,
+        ``fn(a) == fn(b) iff a == b``.
+
+    d: MutableMapping
+        Wrapped mapping
+
+    See Also
+    --------
+    Func
+
+    Examples
+    --------
+    Use any python object as keys of a File, instead of just strings, as long as their
+    str representation is unique:
+
+    >>> from zict import File
+    >>> z = KeyMap(str, File("myfile"))  # doctest: +SKIP
+    >>> z[1] = 10  # doctest: +SKIP
+    """
+
+    fn: Callable[[KT], JT]
+    d: MutableMapping[JT, VT]
+    keymap: dict[KT, JT]
+
+    def __init__(self, fn: Callable[[KT], JT], d: MutableMapping[JT, VT]):
+        super().__init__()
+        self.fn = fn
+        self.d = d
+        self.keymap = {}
+
+    @locked
+    def __setitem__(self, key: KT, value: VT) -> None:
+        j = self.fn(key)
+        self.keymap[key] = j
+        with self.unlock():
+            self.d[j] = value
+        if key not in self.keymap:
+            # Race condition with __delitem__
+            discard(self.d, j)
+
+    def __getitem__(self, key: KT) -> VT:
+        j = self.keymap[key]
+        return self.d[j]
+
+    @locked
+    def __delitem__(self, key: KT) -> None:
+        j = self.keymap.pop(key)
+        del self.d[j]
+
+    def __contains__(self, key: object) -> bool:
+        return key in self.keymap
+
+    def __iter__(self) -> Iterator[KT]:
+        return iter(self.keymap)
+
+    def __len__(self) -> int:
+        return len(self.keymap)
+
+    def flush(self) -> None:
+        flush(self.d)
+
+    def close(self) -> None:
+        close(self.d)

zict/tests/test_keymap.py

@@ -0,0 +1,39 @@
+import pytest
+
+from zict import KeyMap
+from zict.tests import utils_test
+
+
+def test_simple():
+    d = {}
+    z = KeyMap(str, d)
+    z[1] = 10
+    assert d == {"1": 10}
+    assert z.keymap == {1: "1"}
+    assert 1 in z
+    assert 2 not in z
+    assert list(z) == [1]
+    assert len(z) == 1
+    assert z[1] == 10
+    with pytest.raises(KeyError):
+        z[2]
+    del z[1]
+    assert 1 not in z
+    assert 1 not in z.keymap
+
+
+def test_mapping():
+    z = KeyMap(str, {})
+    utils_test.check_mapping(z)
+    utils_test.check_closing(z)
+
+
+@pytest.mark.stress
+@pytest.mark.repeat(utils_test.REPEAT_STRESS_TESTS)
+def test_stress_same_key_threadsafe():
+    d = utils_test.SlowDict(0.001)
+    z = KeyMap(str, d)
+    utils_test.check_same_key_threadsafe(z)
+    assert not z.keymap
+    assert not z.d
+    utils_test.check_mapping(z)