From 42b0f439176db9cd48710d11de341eec6717be2b Mon Sep 17 00:00:00 2001 From: crusaderky Date: Mon, 2 Oct 2023 21:23:57 +0100 Subject: [PATCH] KeyMap --- doc/source/changelog.rst | 1 + doc/source/index.rst | 2 + zict/__init__.py | 1 + zict/func.py | 6 ++- zict/keymap.py | 81 +++++++++++++++++++++++++++++++++++++++ zict/tests/test_keymap.py | 39 +++++++++++++++++++ 6 files changed, 129 insertions(+), 1 deletion(-) create mode 100644 zict/keymap.py create mode 100644 zict/tests/test_keymap.py diff --git a/doc/source/changelog.rst b/doc/source/changelog.rst index 120139e..d6e8dfa 100644 --- a/doc/source/changelog.rst +++ b/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 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..47af8f7 --- /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 + >>> 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) diff --git a/zict/tests/test_keymap.py b/zict/tests/test_keymap.py new file mode 100644 index 0000000..72a843d --- /dev/null +++ b/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)