Add cloud storage support via fsspec (S3, GCS, Azure)

Read and write GeoTIFFs to/from cloud storage using fsspec as the
filesystem abstraction layer. Any URI with a :// scheme (that isn't
http/https) is routed through fsspec, which delegates to the
appropriate backend:

- s3://bucket/key.tif  (requires s3fs)
- gs://bucket/key.tif  (requires gcsfs)
- az://container/blob.tif  (requires adlfs)
- abfs://container/blob.tif  (requires adlfs)
- Any other fsspec-supported scheme (memory://, ftp://, etc.)

Read: _CloudSource uses fsspec.core.url_to_fs() then fs.open()
for full reads and range reads. Falls through to the same TIFF
parsing pipeline as local files.

Write: _write_bytes detects fsspec URIs and writes via fs.open()
instead of the local atomic-rename path (which doesn't apply to
cloud storage).

If fsspec or the backend library isn't installed, a clear ImportError
is raised with install instructions.

5 new tests using fsspec's memory:// filesystem for integration
testing without real cloud credentials.

Author: brendancol

xrspatial/geotiff/_reader.py

@@ -171,10 +171,57 @@ def close(self):
         pass
 
 
+_CLOUD_SCHEMES = ('s3://', 'gs://', 'az://', 'abfs://')
+
+
+def _is_fsspec_uri(path: str) -> bool:
+    """Check if a path is a fsspec-compatible URI (not http/https/local)."""
+    if path.startswith(('http://', 'https://')):
+        return False
+    return '://' in path
+
+
+class _CloudSource:
+    """Cloud storage data source using fsspec.
+
+    Supports S3, GCS, Azure Blob Storage, and any other fsspec backend.
+    Requires the appropriate library (s3fs, gcsfs, adlfs) to be installed.
+    """
+
+    def __init__(self, url: str, **storage_options):
+        try:
+            import fsspec
+        except ImportError:
+            raise ImportError(
+                "fsspec is required to read from cloud storage. "
+                "Install it with: pip install fsspec")
+        self._url = url
+        self._fs, self._path = fsspec.core.url_to_fs(url, **storage_options)
+        self._size = self._fs.size(self._path)
+
+    def read_range(self, start: int, length: int) -> bytes:
+        with self._fs.open(self._path, 'rb') as f:
+            f.seek(start)
+            return f.read(length)
+
+    def read_all(self) -> bytes:
+        with self._fs.open(self._path, 'rb') as f:
+            return f.read()
+
+    @property
+    def size(self) -> int:
+        return self._size
+
+    def close(self):
+        pass
+
+
 def _open_source(source: str):
-    """Open a data source (local file or URL)."""
+    """Open a data source (local file, URL, or cloud path)."""
     if source.startswith(('http://', 'https://')):
         return _HTTPSource(source)
+    if _is_fsspec_uri(source):
+        return _CloudSource(source)
     return _FileSource(source)
 
 
@@ -615,13 +662,14 @@ def read_to_array(source: str, *, window=None, overview_level: int | None = None
     -------
     (np.ndarray, GeoInfo) tuple
     """
-    is_url = source.startswith(('http://', 'https://'))
-
-    if is_url:
+    if source.startswith(('http://', 'https://')):
         return _read_cog_http(source, overview_level=overview_level, band=band)
 
-    # Local file: mmap for zero-copy access
-    src = _FileSource(source)
+    # Local file or cloud storage: read all bytes then parse
+    if _is_fsspec_uri(source):
+        src = _CloudSource(source)
+    else:
+        src = _FileSource(source)
     data = src.read_all()
 
     try:

xrspatial/geotiff/_writer.py

@@ -812,9 +812,33 @@ def write(data: np.ndarray, path: str, *,
         force_bigtiff=bigtiff,
     )
 
-    # Write to a temp file then atomically rename, so concurrent writes to
-    # the same path don't interleave and readers never see partial output.
+    _write_bytes(file_bytes, path)
+
+
+def _is_fsspec_uri(path: str) -> bool:
+    """Check if a path is a fsspec-compatible URI."""
+    if path.startswith(('http://', 'https://')):
+        return False
+    return '://' in path
+
+
+def _write_bytes(file_bytes: bytes, path: str) -> None:
+    """Write bytes to a local file (atomic) or cloud storage (via fsspec)."""
     import os
+
+    if _is_fsspec_uri(path):
+        try:
+            import fsspec
+        except ImportError:
+            raise ImportError(
+                "fsspec is required to write to cloud storage. "
+                "Install it with: pip install fsspec")
+        fs, fspath = fsspec.core.url_to_fs(path)
+        with fs.open(fspath, 'wb') as f:
+            f.write(file_bytes)
+        return
+
+    # Local file: write to temp file then atomically rename
     import tempfile
     dir_name = os.path.dirname(os.path.abspath(path))
     fd, tmp_path = tempfile.mkstemp(dir=dir_name, suffix='.tif.tmp')
@@ -823,7 +847,6 @@ def write(data: np.ndarray, path: str, *,
             f.write(file_bytes)
         os.replace(tmp_path, path)  # atomic on POSIX
     except BaseException:
-        # Clean up the temp file on any failure
         try:
             os.unlink(tmp_path)
         except OSError:

xrspatial/geotiff/tests/test_features.py

@@ -425,6 +425,112 @@ def test_no_crs_no_wkt(self, tmp_path):
 # Big-endian pixel data
 # -----------------------------------------------------------------------
 
+# -----------------------------------------------------------------------
+# Cloud storage (fsspec) support
+# -----------------------------------------------------------------------
+
+class TestCloudStorage:
+
+    def test_cloud_scheme_detection(self):
+        """Cloud URI schemes are detected correctly."""
+        from xrspatial.geotiff._reader import _is_fsspec_uri
+        assert _is_fsspec_uri('s3://bucket/key.tif')
+        assert _is_fsspec_uri('gs://bucket/key.tif')
+        assert _is_fsspec_uri('az://container/blob.tif')
+        assert _is_fsspec_uri('abfs://container/blob.tif')
+        assert _is_fsspec_uri('memory:///test.tif')
+        assert not _is_fsspec_uri('/local/path.tif')
+        assert not _is_fsspec_uri('http://example.com/file.tif')
+        assert not _is_fsspec_uri('relative/path.tif')
+
+    def test_memory_filesystem_read_write(self, tmp_path):
+        """Round-trip through fsspec's in-memory filesystem."""
+        import fsspec
+
+        arr = np.arange(16, dtype=np.float32).reshape(4, 4)
+
+        # Write to memory filesystem via fsspec
+        from xrspatial.geotiff._writer import write, _write_bytes
+        from xrspatial.geotiff._writer import _assemble_tiff, _write_stripped
+        from xrspatial.geotiff._compression import COMPRESSION_NONE
+
+        # First write locally, then copy to memory fs
+        local_path = str(tmp_path / 'test.tif')
+        write(arr, local_path, compression='none', tiled=False)
+
+        with open(local_path, 'rb') as f:
+            tiff_bytes = f.read()
+
+        # Put into fsspec memory filesystem
+        fs = fsspec.filesystem('memory')
+        fs.pipe('/test.tif', tiff_bytes)
+
+        # Read via _CloudSource
+        from xrspatial.geotiff._reader import _CloudSource
+        src = _CloudSource('memory:///test.tif')
+        data = src.read_all()
+        assert len(data) == len(tiff_bytes)
+        assert data == tiff_bytes
+
+        # Range read
+        chunk = src.read_range(0, 8)
+        assert chunk == tiff_bytes[:8]
+
+        # Clean up
+        fs.rm('/test.tif')
+
+    def test_memory_filesystem_full_roundtrip(self, tmp_path):
+        """write_geotiff + read_geotiff through memory:// filesystem."""
+        import fsspec
+
+        arr = np.arange(16, dtype=np.float32).reshape(4, 4)
+
+        # Write locally first, then copy to memory fs
+        local_path = str(tmp_path / 'local.tif')
+        write_geotiff(arr, local_path, compression='deflate')
+        with open(local_path, 'rb') as f:
+            tiff_bytes = f.read()
+
+        fs = fsspec.filesystem('memory')
+        fs.pipe('/roundtrip.tif', tiff_bytes)
+
+        # Read from memory filesystem
+        from xrspatial.geotiff._reader import read_to_array
+        result, geo = read_to_array('memory:///roundtrip.tif')
+        np.testing.assert_array_equal(result, arr)
+
+        fs.rm('/roundtrip.tif')
+
+    def test_writer_cloud_scheme_detection(self):
+        """Writer detects cloud schemes."""
+        from xrspatial.geotiff._writer import _is_fsspec_uri
+        assert _is_fsspec_uri('s3://bucket/key.tif')
+        assert _is_fsspec_uri('gs://bucket/key.tif')
+        assert _is_fsspec_uri('az://container/blob.tif')
+        assert not _is_fsspec_uri('/local/path.tif')
+
+    def test_write_to_memory_filesystem(self, tmp_path):
+        """_write_bytes can write to fsspec memory filesystem."""
+        import fsspec
+        from xrspatial.geotiff._writer import write
+
+        arr = np.arange(16, dtype=np.float32).reshape(4, 4)
+        local_path = str(tmp_path / 'src.tif')
+        write(arr, local_path, compression='none', tiled=False)
+        with open(local_path, 'rb') as f:
+            tiff_bytes = f.read()
+
+        # Write via _write_bytes to memory filesystem
+        from xrspatial.geotiff._writer import _write_bytes
+        _write_bytes(tiff_bytes, 'memory:///written.tif')
+
+        fs = fsspec.filesystem('memory')
+        assert fs.exists('/written.tif')
+        assert fs.cat('/written.tif') == tiff_bytes
+
+        fs.rm('/written.tif')
+
+
 class TestBigEndian:
 
     def test_float32_big_endian(self, tmp_path):