refactoring storage
MODIFY
pangeo_forge_recipes/storage.py
pangeo_forge_recipes/storage.py
import hashlib
+import logging
import os
import re
+import tempfile
import unicodedata
from abc import ABC, abstractmethod
from contextlib import contextmanager
from dataclasses import dataclass
-from typing import Iterator
+from typing import Any, Iterator, Optional, Union
import fsspec
+logger = logging.getLogger(__name__)
+
+# fsspec doesn't provide type hints, so I'm not sure what the write type is for open files
+OpenFileType = Any
+
+
+def _get_url_size(fname):
+ with fsspec.open(fname, mode="rb") as of:
+ size = of.size
+ return size
+
+
+@contextmanager
+def _fsspec_safe_open(fname: str, **kwargs) -> Iterator[OpenFileType]:
+ # workaround for inconsistent behavior of fsspec.open
+ # https://github.com/intake/filesystem_spec/issues/579
+ with fsspec.open(fname, **kwargs) as fp:
+ with fp as fp2:
+ yield fp2
+
+
+def _copy_btw_filesystems(input_opener, output_opener, BLOCK_SIZE=10_000_000):
+ with input_opener as source:
+ with output_opener as target:
+ while True:
+ data = source.read(BLOCK_SIZE)
+ if not data:
+ break
+ target.write(data)
+
class AbstractTarget(ABC):
@abstractmethod
@@ -98,7 +130,58 @@ class FlatFSSpecTarget(FSSpecTarget):
class CacheFSSpecTarget(FlatFSSpecTarget):
"""Alias for FlatFSSpecTarget"""
- pass
+ def cache_file(self, fname: str, **open_kwargs) -> None:
+ # check and see if the file already exists in the cache
+ logger.info(f"Caching file '{fname}'")
+ if self.exists(fname):
+ cached_size = self.size(fname)
+ remote_size = _get_url_size(fname)
+ if cached_size == remote_size:
+ # TODO: add checksumming here
+ logger.info(f"File '{fname}' is already cached")
+ return
+
+ input_opener = _fsspec_safe_open(fname, mode="rb", **open_kwargs)
+ target_opener = self.open(fname, mode="wb")
+ logger.info(f"Coping remote file '{fname}' to cache")
+ _copy_btw_filesystems(input_opener, target_opener)
+
+
+@contextmanager
+def file_opener(
+ fname: str,
+ cache: Optional[CacheFSSpecTarget] = None,
+ copy_to_local: bool = False,
+ **open_kwargs,
+) -> Iterator[Union[OpenFileType, str]]:
+ """
+ Context manager for opening files.
+
+ :param fname: The filename / url to open. Fsspec will inspect the protocol
+ (e.g. http, ftp) and determine the appropriate filesystem type to use.
+ :param cache: A target where the file may have been cached. If none, the file
+ will be opened directly.
+ :param copy_to_local: If True, always copy the file to a local temporary file
+ before opening. In this case, function yields a path name rather than an open file.
+ """
+ if cache is not None:
+ logger.info(f"Opening '{fname}' from cache")
+ opener = cache.open(fname, mode="rb")
+ else:
+ logger.info(f"Opening '{fname}' directly.")
+ opener = _fsspec_safe_open(fname, mode="rb", **open_kwargs)
+ if copy_to_local:
+ _, suffix = os.path.splitext(fname)
+ ntf = tempfile.NamedTemporaryFile(suffix=suffix)
+ tmp_name = ntf.name
+ logger.info(f"Copying '{fname}' to local file '{tmp_name}'")
+ target_opener = open(tmp_name, mode="wb")
+ _copy_btw_filesystems(opener, target_opener)
+ yield tmp_name
+ ntf.close() # cleans up the temporary file
+ else:
+ with opener as fp:
+ yield fp
def _slugify(value: str) -> str:
@@ -116,7 +199,7 @@ class UninitializedTarget(AbstractTarget):
raise UninitializedTargetError
def exists(self, path: str) -> bool:
- raise UninitializedTargetError
+ return False
def rm(self, path: str) -> None:
raise UninitializedTargetError

ADD
tests/test_storage.py
+import pytest
+
+from pangeo_forge_recipes.storage import UninitializedTargetError, file_opener
+
+
+def test_target(tmp_target):
+ mapper = tmp_target.get_mapper()
+ mapper["foo"] = b"bar"
+ with open(tmp_target.root_path + "/foo") as f:
+ res = f.read()
+ assert res == "bar"
+ with pytest.raises(FileNotFoundError):
+ tmp_target.rm("baz")
+ with pytest.raises(FileNotFoundError):
+ with tmp_target.open("baz"):
+ pass
+
+
+def test_uninitialized_target(uninitialized_target):
+ target = uninitialized_target
+ assert not target.exists("foo")
+ with pytest.raises(UninitializedTargetError):
+ target.get_mapper()
+ with pytest.raises(UninitializedTargetError):
+ target.rm("foo")
+ with pytest.raises(UninitializedTargetError):
+ with target.open("foo"):
+ pass
+
+
+def test_cache(tmp_cache):
+ assert not tmp_cache.exists("foo")
+ with tmp_cache.open("foo", mode="w") as f:
+ f.write("bar")
+ assert tmp_cache.exists("foo")
+ assert tmp_cache.size("foo") == 3
+ with tmp_cache.open("foo", mode="r") as f:
+ assert f.read() == "bar"
+ tmp_cache.rm("foo")
+ assert not tmp_cache.exists("foo")
+
+
+@pytest.mark.parametrize("copy_to_local", [False, True])
+@pytest.mark.parametrize("use_cache, cache_first", [(False, False), (True, False), (True, True)])
+def test_file_opener(netcdf_local_paths, tmp_cache, copy_to_local, use_cache, cache_first):
+ all_paths, _ = netcdf_local_paths
+ path = str(all_paths[0])
+
+ cache = tmp_cache if use_cache else None
+ if cache_first:
+ cache.cache_file(path)
+ opener = file_opener(path, cache, copy_to_local=copy_to_local)
+ if use_cache and not cache_first:
+ with pytest.raises(FileNotFoundError):
+ with opener as fp:
+ pass
+ else:
+ with opener as fp:
+ if copy_to_local:
+ assert isinstance(fp, str)
+ with open(fp, mode="rb") as fp2:
+ _ = fp2.read()
+ else:
+ _ = fp.read()

refactoring storage