diff --git a/fsspec/caching.py b/fsspec/caching.py index b6d92587e..828a67061 100644 --- a/fsspec/caching.py +++ b/fsspec/caching.py @@ -767,14 +767,39 @@ def _read_cache(self, start, end, start_block_number, end_block_number): caches = { - "none": BaseCache, + # one custom case None: BaseCache, - "mmap": MMapCache, - "bytes": BytesCache, - "readahead": ReadAheadCache, - "block": BlockCache, - "first": FirstChunkCache, - "all": AllBytes, - "parts": KnownPartsOfAFile, - "background": BackgroundBlockCache, } + + +def register_cache(cls, clobber=False): + """'Register' cache implementation. + + Parameters + ---------- + clobber: bool, optional + If set to True (default is False) - allow to overwrite existing + entry. + + Raises + ------ + ValueError + """ + name = cls.name + if not clobber and name in caches: + raise ValueError(f"Cache with name {name!r} is already known: {caches[name]}") + caches[name] = cls + + +for c in ( + BaseCache, + MMapCache, + BytesCache, + ReadAheadCache, + BlockCache, + FirstChunkCache, + AllBytes, + KnownPartsOfAFile, + BackgroundBlockCache, +): + register_cache(c) diff --git a/fsspec/tests/test_caches.py b/fsspec/tests/test_caches.py index c655e2ae7..65c80cdb4 100644 --- a/fsspec/tests/test_caches.py +++ b/fsspec/tests/test_caches.py @@ -3,7 +3,7 @@ import pytest -from fsspec.caching import BlockCache, FirstChunkCache, caches +from fsspec.caching import BlockCache, FirstChunkCache, caches, register_cache def test_cache_getitem(Cache_imp): @@ -138,3 +138,10 @@ def wrapped(*a, **kw): f.read(1) time.sleep(0.1) # second block is loading assert len(thread_ids) == 2 + + +def test_register_cache(): + # just test that we have them populated and fail to re-add again unless overload + with pytest.raises(ValueError): + register_cache(BlockCache) + register_cache(BlockCache, clobber=True)