avoid duplication and allow for formalized addition of caches (#1227)

Co-authored-by: Martin Durant <[email protected]>

Author: yarikoptic

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)

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)