Typesafe accessor for Storage

bluesky · Dec 13, 2024 · d838f77 · d838f77
1 parent 9da38f1
commit d838f77
Show file tree

Hide file tree

Showing 8 changed files with 16 additions and 9 deletions.
diff --git a/tiled/adapters/arrow.py b/tiled/adapters/arrow.py
@@ -78,7 +78,7 @@ def init_storage(
         The list of assets.
         """
         data_source = copy.deepcopy(data_source)  # Do not mutate caller input.
-        data_uri = str(storage.filesystem) + "".join(
+        data_uri = str(storage.get("filesystem")) + "".join(
             f"/{quote_plus(segment)}" for segment in path_parts
         )
         directory = path_from_uri(data_uri)

diff --git a/tiled/adapters/awkward_buffers.py b/tiled/adapters/awkward_buffers.py
@@ -43,7 +43,7 @@ def init_storage(
 
         """
         data_source = copy.deepcopy(data_source)  # Do not mutate caller input.
-        data_uri = str(storage.filesystem) + "".join(
+        data_uri = str(storage.get("filesystem")) + "".join(
             f"/{quote_plus(segment)}" for segment in path_parts
         )
         directory: Path = path_from_uri(data_uri)

diff --git a/tiled/adapters/csv.py b/tiled/adapters/csv.py
@@ -145,7 +145,7 @@ def init_storage(
 
         """
         data_source = copy.deepcopy(data_source)  # Do not mutate caller input.
-        data_uri = str(storage.filesystem) + "".join(
+        data_uri = str(storage.get("filesystem")) + "".join(
             f"/{quote_plus(segment)}" for segment in path_parts
         )
         directory = path_from_uri(data_uri)

diff --git a/tiled/adapters/parquet.py b/tiled/adapters/parquet.py
@@ -90,7 +90,7 @@ def init_storage(
 
         """
         data_source = copy.deepcopy(data_source)  # Do not mutate caller input.
-        data_uri = str(storage.filesystem) + "".join(
+        data_uri = str(storage.get("filesystem")) + "".join(
             f"/{quote_plus(segment)}" for segment in path_parts
         )
         directory = path_from_uri(data_uri)

diff --git a/tiled/adapters/sparse_blocks_parquet.py b/tiled/adapters/sparse_blocks_parquet.py
@@ -89,7 +89,7 @@ def init_storage(
 
         """
         data_source = copy.deepcopy(data_source)  # Do not mutate caller input.
-        data_uri = str(storage.filesystem) + "".join(
+        data_uri = str(storage.get("filesystem")) + "".join(
             f"/{quote_plus(segment)}" for segment in path_parts
         )
         directory = path_from_uri(data_uri)

diff --git a/tiled/adapters/sql.py b/tiled/adapters/sql.py
@@ -81,6 +81,7 @@ def init_storage(
     ) -> DataSource[TableStructure]:
         """
         Class to initialize the list of assets for given uri. In SQL Adapter we hve  single partition.
+
         Parameters
         ----------
         data_uri : the uri of the data
@@ -96,7 +97,7 @@ def init_storage(
         default_table_name = ...  # based on hash of Arrow schema
         data_source.parameters.setdefault("table_name", default_table_name)
         data_source.parameters["dataset_id"] = uuid.uuid4().int
-        data_uri = storage.sql  # TODO scrub credentials
+        data_uri = storage.get("sql")  # TODO scrub credentials
         data_source.assets.append(
             Asset(
                 data_uri=data_uri,

diff --git a/tiled/adapters/zarr.py b/tiled/adapters/zarr.py
@@ -76,7 +76,7 @@ def init_storage(
 
         """
         data_source = copy.deepcopy(data_source)  # Do not mutate caller input.
-        data_uri = str(storage.filesystem) + "".join(
+        data_uri = str(storage.get("filesystem")) + "".join(
             f"/{quote_plus(segment)}" for segment in path_parts
         )
         # Zarr requires evenly-sized chunks within each dimension.
@@ -85,9 +85,9 @@ def init_storage(
         shape = tuple(dim[0] * len(dim) for dim in data_source.structure.chunks)
         directory = path_from_uri(data_uri)
         directory.mkdir(parents=True, exist_ok=True)
-        storage = zarr.storage.DirectoryStore(str(directory))
+        store = zarr.storage.DirectoryStore(str(directory))
         zarr.storage.init_array(
-            storage,
+            store,
             shape=shape,
             chunks=zarr_chunks,
             dtype=data_source.structure.data_type.to_numpy_dtype(),

diff --git a/tiled/structures/data_source.py b/tiled/structures/data_source.py
@@ -45,3 +45,9 @@ def from_json(cls, d):
 class Storage:
     filesystem: Optional[str]
     sql: Optional[str]
+
+    def get(self, storage: str) -> str:
+        uri = getattr(self, storage)
+        if isinstance(uri, str):
+            return uri
+        raise TypeError(f"{storage} is not set")