diff --git a/cachier/base_core.py b/cachier/base_core.py index 99b1d9a..6a9bb01 100644 --- a/cachier/base_core.py +++ b/cachier/base_core.py @@ -28,7 +28,7 @@ class _BaseCore(): if such a mapping exists.""" @abc.abstractmethod - def get_entry(self, args, kwds): + def get_entry(self, args, kwds, hash_params): """Returns the result mapped to the given arguments in this core's cache, if such a mapping exists.""" diff --git a/cachier/core.py b/cachier/core.py index 793fe38..a909915 100644 --- a/cachier/core.py +++ b/cachier/core.py @@ -79,6 +79,7 @@ def cachier( pickle_reload=True, mongetter=None, cache_dir=None, + hash_params=None, ): """A persistent, stale-free memoization decorator. @@ -112,6 +113,11 @@ def cachier( A fully qualified path to a file directory to be used for cache files. The running process must have running permissions to this folder. If not provided, a default directory at `~/.cachier/` is used. + hash_params : callable, optional + A callable that takes args and kwargs from main function and returns + a hash key of these params. If unset, default transformation is + applied. It is valuable and works as workaround in scenarios + that positional and keyword arguments are not hashable. """ # print('Inside the wrapper maker') # print('mongetter={}'.format(mongetter)) @@ -142,7 +148,7 @@ def cachier( _print = print if ignore_cache: return func(*args, **kwds) - key, entry = core.get_entry(args, kwds) + key, entry = core.get_entry(args, kwds, hash_params) if overwrite_cache: return _calc_entry(core, key, func, args, kwds) if entry is not None: # pylint: disable=R0101 diff --git a/cachier/mongo_core.py b/cachier/mongo_core.py index 2d2077d..e05b6da 100644 --- a/cachier/mongo_core.py +++ b/cachier/mongo_core.py @@ -79,8 +79,8 @@ class _MongoCore(_BaseCore): return key, entry return key, None - def get_entry(self, args, kwds): - key = pickle.dumps(args + tuple(sorted(kwds.items()))) + def get_entry(self, args, kwds, hash_params): + key = pickle.dumps(args + tuple(sorted(kwds.items())) if hash_params is None else hash_params(args, kwds)) return self.get_entry_by_key(key) def set_entry(self, key, func_res): diff --git a/cachier/pickle_core.py b/cachier/pickle_core.py index e6bdea0..fed5bb4 100644 --- a/cachier/pickle_core.py +++ b/cachier/pickle_core.py @@ -146,8 +146,8 @@ class _PickleCore(_BaseCore): self._reload_cache() return key, self._get_cache().get(key, None) - def get_entry(self, args, kwds): - key = args + tuple(sorted(kwds.items())) + def get_entry(self, args, kwds, hash_params): + key = args + tuple(sorted(kwds.items())) if hash_params is None else hash_params(args, kwds) # print('key type={}, key={}'.format(type(key), key)) return self.get_entry_by_key(key) diff --git a/setup.py b/setup.py index c438286..1bdd079 100644 --- a/setup.py +++ b/setup.py @@ -15,7 +15,7 @@ except ImportError: import versioneer -TEST_REQUIRES = ['pytest', 'coverage', 'pytest-cov', 'pymongo'] +TEST_REQUIRES = ['pytest', 'coverage', 'pytest-cov', 'pymongo', 'pandas'] README_RST = '' with open('README.rst') as f: diff --git a/tests/test_mongo_core.py b/tests/test_mongo_core.py index e62e8d2..b6fd845 100644 --- a/tests/test_mongo_core.py +++ b/tests/test_mongo_core.py @@ -14,6 +14,8 @@ except ImportError: # python 2 import pytest import pymongo +import hashlib +import pandas as pd from pymongo.mongo_client import MongoClient from pymongo.errors import OperationFailure @@ -195,7 +197,7 @@ def test_stalled_mongo_db_cache(): def test_stalled_mong_db_core(monkeypatch): - def mock_get_entry(self, args, kwargs): # skipcq: PYL-R0201, PYL-W0613 + def mock_get_entry(self, args, kwargs, hash_params): # skipcq: PYL-R0201, PYL-W0613 return "key", {'being_calculated': True} def mock_get_entry_by_key(self, key): # skipcq: PYL-R0201, PYL-W0613 @@ -213,7 +215,7 @@ def test_stalled_mong_db_core(monkeypatch): res = _stalled_func() assert res == 1 - def mock_get_entry_2(self, args, kwargs): # skipcq: PYL-W0613 + def mock_get_entry_2(self, args, kwargs, hash_params): # skipcq: PYL-W0613 entry = { 'being_calculated': True, "value": 1, @@ -233,3 +235,35 @@ def test_stalled_mong_db_core(monkeypatch): res = _stalled_func_2() assert res == 2 + + +def test_callable_hash_param(): + + def _hash_params(args, kwargs): + def _hash(obj): + if isinstance(obj, pd.core.frame.DataFrame): + return hashlib.sha256(pd.util.hash_pandas_object(obj).values.tobytes()).hexdigest() + return obj + + k_args = tuple(map(_hash, args)) + k_kwargs = tuple(sorted({k: _hash(v) for k, v in kwargs.items()}.items())) + return k_args + k_kwargs + + @cachier(mongetter=_test_mongetter, hash_params=_hash_params) + def _params_with_dataframe(*args, **kwargs): + """Some function.""" + return random() + + _params_with_dataframe.clear_cache() + + df_a = pd.DataFrame.from_dict(dict(a=[0], b=[2], c=[3])) + df_b = pd.DataFrame.from_dict(dict(a=[0], b=[2], c=[3])) + value_a = _params_with_dataframe(df_a, 1) + value_b = _params_with_dataframe(df_b, 1) + + assert value_a == value_b # same content --> same key + + value_a = _params_with_dataframe(1, df=df_a) + value_b = _params_with_dataframe(1, df=df_b) + + assert value_a == value_b # same content --> same key diff --git a/tests/test_pickle_core.py b/tests/test_pickle_core.py index 0fef615..abf2591 100644 --- a/tests/test_pickle_core.py +++ b/tests/test_pickle_core.py @@ -24,6 +24,9 @@ try: except ImportError: # python 2 import Queue as queue +import hashlib +import pandas as pd + from cachier import cachier from cachier.pickle_core import DEF_CACHIER_DIR @@ -392,3 +395,35 @@ def test_pickle_core_custom_cache_dir(): assert end - start < 1 _takes_5_seconds_custom_dir.clear_cache() assert _takes_5_seconds_custom_dir.cache_dpath() == EXPANDED_CUSTOM_DIR + + +def test_callable_hash_param(): + + def _hash_params(args, kwargs): + def _hash(obj): + if isinstance(obj, pd.core.frame.DataFrame): + return hashlib.sha256(pd.util.hash_pandas_object(obj).values.tobytes()).hexdigest() + return obj + + k_args = tuple(map(_hash, args)) + k_kwargs = tuple(sorted({k: _hash(v) for k, v in kwargs.items()}.items())) + return k_args + k_kwargs + + @cachier(hash_params=_hash_params) + def _params_with_dataframe(*args, **kwargs): + """Some function.""" + return random() + + _params_with_dataframe.clear_cache() + + df_a = pd.DataFrame.from_dict(dict(a=[0], b=[2], c=[3])) + df_b = pd.DataFrame.from_dict(dict(a=[0], b=[2], c=[3])) + value_a = _params_with_dataframe(df_a, 1) + value_b = _params_with_dataframe(df_b, 1) + + assert value_a == value_b # same content --> same key + + value_a = _params_with_dataframe(1, df=df_a) + value_b = _params_with_dataframe(1, df=df_b) + + assert value_a == value_b # same content --> same key