From 99faca94469102979fc68d6fa86b5d54c4b786e6 Mon Sep 17 00:00:00 2001
From: crusaderky <crusaderky@gmail.com>
Date: Mon, 2 Oct 2023 21:18:05 +0100
Subject: [PATCH] KeyMap

---
 doc/source/index.rst      |  2 +
 zict/__init__.py          |  1 +
 zict/func.py              |  6 ++-
 zict/keymap.py            | 81 +++++++++++++++++++++++++++++++++++++++
 zict/tests/test_keymap.py | 36 +++++++++++++++++
 5 files changed, 125 insertions(+), 1 deletion(-)
 create mode 100644 zict/keymap.py
 create mode 100644 zict/tests/test_keymap.py

diff --git a/doc/source/index.rst b/doc/source/index.rst
index d37d8c3..6769b79 100644
--- a/doc/source/index.rst
+++ b/doc/source/index.rst
@@ -55,6 +55,8 @@ API
    :members:
 .. autoclass:: Func
    :members:
+.. autoclass:: KeyMap
+   :members:
 .. autoclass:: LMDB
    :members:
 .. autoclass:: LRU
diff --git a/zict/__init__.py b/zict/__init__.py
index 11527f1..49ee47e 100644
--- a/zict/__init__.py
+++ b/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
diff --git a/zict/func.py b/zict/func.py
index b45cfa0..3fb6e8c 100644
--- a/zict/func.py
+++ b/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):
diff --git a/zict/keymap.py b/zict/keymap.py
new file mode 100644
index 0000000..ba7b7d7
--- /dev/null
+++ b/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
+    >>> f = KeyMap(str, File("myfile"))  # doctest: +SKIP
+    >>> f[1] = 10
+    """
+
+    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)
diff --git a/zict/tests/test_keymap.py b/zict/tests/test_keymap.py
new file mode 100644
index 0000000..b3fdb0d
--- /dev/null
+++ b/zict/tests/test_keymap.py
@@ -0,0 +1,36 @@
+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():
+    z = KeyMap(str, {})
+    utils_test.check_same_key_threadsafe(z)
+    utils_test.check_mapping(z)