Fix storage options (#83)

* Fix storage options * Update method docs Co-authored-by: Alp Arıbal <[email protected]> * Update init comment Co-authored-by: Alp Arıbal <[email protected]> * Another version bump Co-authored-by: Alp Arıbal <[email protected]>
merantix-momentum · Aug 31, 2022 · 721e1f0 · 721e1f0
1 parent 1331593
commit 721e1f0
Show file tree

Hide file tree

Showing 3 changed files with 14 additions and 6 deletions.
diff --git a/squirrel/__init__.py b/squirrel/__init__.py
@@ -1 +1 @@
-__version__ = "0.17.3"
+__version__ = "0.17.4"
diff --git a/squirrel/driver/store_driver.py b/squirrel/driver/store_driver.py
@@ -1,6 +1,6 @@
 from __future__ import annotations
 
-from typing import TYPE_CHECKING, Any, Iterable
+from typing import TYPE_CHECKING, Any, Iterable, Dict, Optional
 
 from squirrel.driver.driver import MapDriver
 from squirrel.serialization import SquirrelSerializer
@@ -19,17 +19,21 @@ class StoreDriver(MapDriver):
 
     name = "store_driver"
 
-    def __init__(self, url: str, serializer: SquirrelSerializer, **kwargs) -> None:
+    def __init__(
+        self, url: str, serializer: SquirrelSerializer, storage_options: Optional[Dict[str, Any]] = None, **kwargs
+    ) -> None:
         """Initializes StoreDriver.
 
         Args:
             url (str): the url of the store
             serializer (SquirrelSerializer): serializer to be passed to SquirrelStore
+            storage_options (Optional[Dict[str, Any]]): a dict with keyword arguments to be passed to store initializer
             **kwargs: Keyword arguments to pass to the super class initializer.
         """
         super().__init__(**kwargs)
         self.url = url
         self.serializer = serializer
+        self.storage_options = storage_options if storage_options is not None else {}
         self._store = None
 
     def get_iter(self, flatten: bool = True, **kwargs) -> Composable:
@@ -83,5 +87,5 @@ def keys(self, **kwargs) -> Iterable:
     def store(self) -> AbstractStore:
         """Store that is used by the driver."""
         if self._store is None:
-            self._store = SquirrelStore(url=self.url, serializer=self.serializer)
+            self._store = SquirrelStore(url=self.url, serializer=self.serializer, **self.storage_options)
         return self._store
diff --git a/squirrel/store/filesystem.py b/squirrel/store/filesystem.py
@@ -87,11 +87,15 @@ def keys(self, nested: bool = True, **kwargs) -> t.Iterator[str]:
             nested (bool): Whether to return paths that are not direct children of the root directory. If True, all
                 paths in the store will be yielded. Otherwise, only the top-level paths (i.e. direct children of the
                 root path) will be yielded. This option is passed to FilePathGenerator initializer. Defaults to True.
-            **kwargs: Other keyword arguments passed to the FilePathGenerator initializer.
+            **kwargs: Other keyword arguments passed to the FilePathGenerator initializer. If a key is present in
+                both `kwargs` and `self.storage_options`, the value from `kwargs` will be used.
 
         Yields:
             (str) Paths to files and directories in the store relative to the root directory.
         """
-        fp_gen = FilePathGenerator(url=self.url, nested=nested, **kwargs)
+        storage_options = self.storage_options.copy()
+        storage_options.update(kwargs)
+
+        fp_gen = FilePathGenerator(url=self.url, nested=nested, **storage_options)
         for path in fp_gen:
             yield osp.relpath(path, start=self.url)