stac reader

.circleci/config.yml

@@ -29,7 +29,7 @@ jobs:
       # https://nose.readthedocs.io
       - run:
           name: run tests
-          command: python setup.py pytest
+          command: pip install -e .[dev] && pytest -v
 
       - codecov/upload:
-          file: "coverage.xml"
+          file: "coverage.xml"

aiocogeo/__init__.py

@@ -1 +1,2 @@
-from .cog import COGReader
+from .cog import COGReader, CompositeReader
+from .stac import STACReader

aiocogeo/cog.py

@@ -1,8 +1,9 @@
+import abc
 import asyncio
 from dataclasses import dataclass, field
 import logging
 import math
-from typing import Any, Dict, List, Optional, Tuple, Union
+from typing import Any, Callable, Dict, List, Optional, Tuple, Union
 from urllib.parse import urljoin
 import uuid
 
@@ -23,7 +24,44 @@
 
 
 @dataclass
-class COGReader(PartialReadInterface):
+class ReaderMixin(abc.ABC):
+
+    async def __aenter__(self):
+        return self
+
+    async def __aexit__(self, exc_type, exc_val, exc_tb):
+        ...
+
+    @abc.abstractmethod
+    async def get_tile(self, x: int, y: int, z: int) -> Union[np.ndarray, List[np.ndarray]]:
+        ...
+
+    @abc.abstractmethod
+    async def read(
+        self,
+        bounds: Tuple[float, float, float, float],
+        shape: Tuple[int, int],
+        resample_method: int = Image.NEAREST,
+    ) -> Union[Union[np.ndarray, np.ma.masked_array], List[Union[np.ndarray, np.ma.masked_array]]]:
+        ...
+
+    @abc.abstractmethod
+    async def point(self, x: Union[float, int], y: Union[float, int]) -> Union[Union[np.ndarray, np.ma.masked_array], List[Union[np.ndarray, np.ma.masked_array]]]:
+        ...
+
+    @abc.abstractmethod
+    async def preview(
+        self,
+        max_size: int = 1024,
+        height: Optional[int] = None,
+        width: Optional[int] = None,
+        resample_method: int = Image.NEAREST
+    ) -> Union[Union[np.ndarray, np.ma.masked_array], List[Union[np.ndarray, np.ma.masked_array]]]:
+        ...
+
+
+@dataclass
+class COGReader(ReaderMixin, PartialReadInterface):
     filepath: str
     ifds: Optional[List[ImageIFD]] = field(default_factory=lambda: [])
     mask_ifds: Optional[List[MaskIFD]] = field(default_factory=lambda: [])
@@ -323,3 +361,69 @@ def create_tile_matrix_set(self, identifier: str = None) -> Dict[str, Any]:
             "tileMatrix": list(reversed(matrices))
         }
         return tms
+
+
+FilterType = Callable[[COGReader], Any]
+MapType = Callable[[COGReader], Any]
+ReduceType = Callable[[List[Union[np.ndarray, np.ma.masked_array]]], Any]
+
+@dataclass
+class CompositeReader(ReaderMixin):
+    readers: Optional[List[COGReader]] = field(default_factory=list)
+    filter: FilterType = lambda a: a
+    default_reducer: ReduceType = lambda r: r
+
+    def __iter__(self):
+        return iter(self.readers)
+
+    async def __aenter__(self):
+        return self
+
+    async def __aexit__(self, exc_type, exc_val, exc_tb):
+        pass
+
+    async def map(self, func: MapType) -> List[Any]:
+        futs = [func(reader) for reader in filter(self.filter, self.readers)]
+        return await asyncio.gather(*futs)
+
+    async def get_tile(self, x: int, y: int, z: int, reducer: Optional[ReduceType] = None) -> List[np.ndarray]:
+        """Fetch a tile from all readers"""
+        tiles = await self.map(
+            func=lambda r: r.get_tile(x, y, z),
+        )
+        reducer = reducer or self.default_reducer
+        return reducer(tiles)
+
+    async def read(
+        self,
+        bounds: Tuple[float, float, float, float],
+        shape: Tuple[int, int],
+        resample_method: int = Image.NEAREST,
+        reducer: Optional[ReduceType] = None
+    ):
+        reads = await self.map(
+            func=lambda r: r.read(bounds, shape, resample_method)
+        )
+        reducer = reducer or self.default_reducer
+        return reducer(reads)
+
+    async def point(self, x: Union[float, int], y: Union[float, int], reducer: Optional[ReduceType] = None) -> List[Union[np.ndarray, np.ma.masked_array]]:
+        points = await self.map(
+            func=lambda r: r.point(x, y)
+        )
+        reducer = reducer or self.default_reducer
+        return reducer(points)
+
+    async def preview(
+        self,
+        max_size: int = 1024,
+        height: Optional[int] = None,
+        width: Optional[int] = None,
+        resample_method: int = Image.NEAREST,
+        reducer: Optional[ReduceType] = None,
+    ) -> List[Union[np.ndarray, np.ma.masked_array]]:
+        previews = await self.map(
+            func=lambda r: r.preview(max_size, height, width, resample_method)
+        )
+        reducer = reducer or self.default_reducer
+        return reducer(previews)

aiocogeo/errors.py

@@ -14,3 +14,8 @@ class InvalidTiffError(CogReadError):
 @dataclass
 class TileNotFoundError(CogReadError):
     ...
+
+
+@dataclass
+class MissingAssets(CogReadError):
+    ...

aiocogeo/filesystems.py

@@ -1,6 +1,7 @@
 import abc
 import asyncio
 from dataclasses import dataclass, field
+import json
 import logging
 import time
 from typing import Any, Callable, Dict, Union
@@ -86,6 +87,10 @@ async def range_request(self, start: int, offset: int) -> bytes:
         """
         return await self._range_request(start, offset)
 
+    @abc.abstractmethod
+    async def request_json(self):
+        ...
+
     @abc.abstractmethod
     async def _range_request(self, start: int, offset: int) -> bytes:
         """Perform a range request"""
@@ -149,6 +154,17 @@ async def _range_request(self, start: int, offset: int) -> bytes:
             raise FileNotFoundError(f"File not found: {self.filepath}") from e
         return data
 
+    async def request_json(self) -> Dict:
+        try:
+            async with self.session.get(self.filepath) as resp:
+                resp.raise_for_status()
+                data = await resp.json()
+        except (aiohttp.ClientError, aiohttp.ClientResponseError) as e:
+            await self._close()
+            raise FileNotFoundError(f"File not found: {self.filepath}") from e
+        return data
+
+
     async def _close(self) -> None:
         if 'session' not in self.kwargs:
             await self.session.close()
@@ -173,10 +189,11 @@ async def _on_request_start(self, session, trace_config_ctx, params):
     async def _on_request_end(self, session, trace_config_ctx, params):
         if params.response.status < 400:
             elapsed = round(asyncio.get_event_loop().time() - trace_config_ctx.start, 3)
-            content_range = params.response.headers['Content-Range']
+            content_range = params.response.headers.get('Content-Range')
             self._total_bytes_requested += int(params.response.headers["Content-Length"])
             self._total_requests += 1
-            self._requested_ranges.append(tuple([int(v) for v in content_range.split(' ')[-1].split('/')[0].split('-')]))
+            if content_range:
+                self._requested_ranges.append(tuple([int(v) for v in content_range.split(' ')[-1].split('/')[0].split('-')]))
             if config.VERBOSE_LOGS:
                 debug_statement = [f"\n < HTTP/{session.version.major}.{session.version.minor}"]
                 debug_statement += [f"\n < {k}: {v}" for (k, v) in params.response.headers.items()]
@@ -200,6 +217,9 @@ async def _range_request(self, start: int, offset: int) -> bytes:
         logger.debug(f" FINISHED REQUEST in {elapsed} seconds: <STATUS 206> ({start}-{start+offset})")
         return data
 
+    async def request_json(self):
+        return json.load(self.file)
+
     async def _close(self) -> None:
         await self.file.close()
 
@@ -232,6 +252,27 @@ async def _range_request(self, start: int, offset: int) -> bytes:
         data = await req['Body'].read()
         return data
 
+    async def request_json(self):
+        kwargs = {}
+        if config.AWS_REQUEST_PAYER:
+            kwargs['RequestPayer'] = config.AWS_REQUEST_PAYER
+        begin = time.time()
+        try:
+            req = await self.object.get()
+        except botocore.exceptions.ClientError as e:
+            await self._close()
+            raise FileNotFoundError(f"File not found: {self.filepath}") from e
+        elapsed = time.time() - begin
+        content_range = req['ResponseMetadata']['HTTPHeaders']['content-range']
+        if not config.VERBOSE_LOGS:
+            status = req['ResponseMetadata']['HTTPStatusCode']
+            logger.debug(f" FINISHED REQUEST in {elapsed} seconds: <STATUS {status}> ({content_range})")
+        self._total_bytes_requested += int(req['ResponseMetadata']['HTTPHeaders']['content-length'])
+        self._total_requests += 1
+        self._requested_ranges.append(tuple([int(v) for v in content_range.split(' ')[-1].split('/')[0].split('-')]))
+        data = json.loads(await req['Body'].read().decode('utf-8'))
+        return data
+
     async def _close(self) -> None:
         await self.resource.__aexit__('', '', '')
 

aiocogeo/stac.py

@@ -0,0 +1,50 @@
+import asyncio
+from dataclasses import dataclass, field
+from typing import Dict, Optional, Set
+
+from stac_pydantic.shared import Asset, MimeTypes
+
+from .cog import COGReader, CompositeReader
+from .filesystems import Filesystem
+from .errors import MissingAssets
+
+
+@dataclass
+class AssetReader(COGReader):
+    asset: Asset = Asset
+
+
+@dataclass
+class STACReader(CompositeReader):
+    filepath: Optional[str] = None
+    include_types: Set[MimeTypes] = field(default_factory=lambda: {MimeTypes.cog})
+
+    kwargs: Optional[Dict] = field(default_factory=dict)
+
+
+    async def __aenter__(self):
+        async with Filesystem.create_from_filepath(self.filepath, **self.kwargs) as file_reader:
+            self._file_reader = file_reader
+            item = await file_reader.request_json()
+
+        # Create a reader for each asset with a COG mime type
+        reader_futs = []
+        for asset in item["assets"]:
+            if item["assets"][asset]["type"] in self.include_types:
+                reader = AssetReader(
+                    filepath=item["assets"][asset]["href"],
+                    asset=Asset(name=asset, **item['assets'][asset])
+                )
+                reader_futs.append(reader)
+
+        if not reader_futs:
+            raise MissingAssets(f"No assets found of type {self.include_types}")
+
+        reader_futs = map(lambda r: r.__aenter__(), filter(self.filter, reader_futs))
+        self.readers = await asyncio.gather(*reader_futs)
+        return self
+
+    async def __aexit__(self, exc_type, exc_val, exc_tb):
+        await self._file_reader._close()
+        for reader in self.readers:
+            await reader._file_reader._close()