From 2d761f0aea3ccb130ab31a4572f5ef70fe9638d0 Mon Sep 17 00:00:00 2001 From: Leo Schick Date: Tue, 19 Sep 2023 15:34:37 +0200 Subject: [PATCH] fix match issue of mara_storage.config.storages --- mara_pipelines/config.py | 14 +++++++++++--- 1 file changed, 11 insertions(+), 3 deletions(-) diff --git a/mara_pipelines/config.py b/mara_pipelines/config.py index ecb677f..d5e1820 100644 --- a/mara_pipelines/config.py +++ b/mara_pipelines/config.py @@ -37,9 +37,17 @@ def default_storage_alias() -> str: """The alias of the storage that should be used when not specified otherwise""" return 'data' -@patch(mara_storage.config.storages) -def storages() -> {str: mara_storage.storages.Storage}: - return {'data': mara_storage.storages.LocalStorage(base_path=pathlib.Path(data_dir()))} +patch(mara_storage.config.storages)(lambda: {'data': mara_storage.storages.LocalStorage(base_path=pathlib.Path(data_dir()))}) + +# Need to patch this function to fix lru_cache issues. mara_storage.storages.storage.cache_clear() does not fix issues here +@functools.lru_cache(maxsize=None) +@patch(mara_storage.storages.storage) +def storage(alias) -> mara_storage.storages.Storage: + """Returns a storage configuration by alias""" + storages = mara_storage.config.storages() + if alias not in storages: + raise KeyError(f'storage alias "{alias}" not configured') + return storages[alias] def default_task_max_retries():