drb-chunk 0.2.0__py3-none-any.whl

drb/chunk/__init__.py

@@ -0,0 +1,16 @@
+from . import _version
+from drb.chunk.exceptions import DrbChunkError, DrbSelectionError
+from drb.chunk.selection import Selection, parse_selection
+from drb.chunk.tiling import TilingScheme, RegularGrid
+from drb.chunk.model import ChunkArray, ChunkRef, ChunkManifest
+from drb.chunk.chunk import Chunk
+from drb.chunk.core import ChunkAddon
+from drb.chunk.interop import to_kerchunk
+
+__version__ = _version.get_versions()["version"]
+
+__all__ = [
+    "Chunk", "ChunkAddon", "ChunkArray", "ChunkRef", "ChunkManifest",
+    "Selection", "parse_selection", "TilingScheme", "RegularGrid",
+    "to_kerchunk", "DrbChunkError", "DrbSelectionError", "__version__",
+]

drb/chunk/_version.py

@@ -0,0 +1,21 @@
+
+# This file was generated by 'versioneer.py' (0.29) from
+# revision-control system data, or from the parent directory name of an
+# unpacked source archive. Distribution tarballs contain a pre-generated copy
+# of this file.
+
+import json
+
+version_json = '''
+{
+ "date": "2026-06-30T04:50:46+0000",
+ "dirty": false,
+ "error": null,
+ "full-revisionid": "c2cf2e26834f78f38f0a6a5162142494c001af69",
+ "version": "0.2.0"
+}
+'''  # END VERSION_JSON
+
+
+def get_versions():
+    return json.loads(version_json)

drb/chunk/chunk.py

@@ -0,0 +1,140 @@
+# drb/chunk/chunk.py
+from typing import Any, Iterable, Optional, Tuple
+
+import numpy as np
+
+from drb.core.node import DrbNode
+
+from drb.chunk.exceptions import DrbChunkError
+from drb.chunk.model import ChunkArray, ChunkRef, ChunkManifest
+from drb.chunk.readers import select_reader
+from drb.chunk.selection import Selection
+
+
+class Chunk:
+    def __init__(
+            self, name: str, array: ChunkArray, node: DrbNode,
+            manifest: ChunkManifest, *, selection: Optional[Selection] = None,
+            reader: Optional[str] = None, topic_uri: Optional[str] = None,
+            addon_data: Optional[dict] = None):
+        self.name = name
+        self.array = array
+        self.node = node
+        self._manifest = manifest
+        self._selection = selection
+        self._reader = reader
+        self._topic_uri = topic_uri
+        self.addon_data = addon_data or {}
+
+    @property
+    def grid_shape(self) -> Tuple[int, ...]:
+        return self.array.grid_shape
+
+    @property
+    def manifest(self) -> ChunkManifest:
+        return self._manifest
+
+    def tiles(self) -> Iterable[Tuple[int, ...]]:
+        return self._manifest.keys()
+
+    def tile(self, key: Tuple[int, ...]) -> ChunkRef:
+        return self._manifest.ref(key)
+
+    def select(self, selection: Selection) -> "Chunk":
+        """Lazy narrowing: resolve keys via the scheme, subset the manifest,
+        read nothing."""
+        resolved = self.array.scheme.resolve(selection, self.array)
+        return Chunk(
+            name=self.name, array=self.array, node=self.node,
+            manifest=self._manifest.subset(resolved),
+            selection=selection, reader=self._reader,
+            topic_uri=self._topic_uri, addon_data=self.addon_data,
+        )
+
+    def __getitem__(self, selection: Selection) -> "Chunk":
+        return self.select(selection)
+
+    def get_impl(self, impl: type, **kwargs) -> Any:
+        if impl is np.ndarray:
+            return self._materialise_numpy()
+        try:
+            import xarray
+        except ImportError:                 # pragma: no cover
+            xarray = None
+        if xarray is not None and impl is xarray.DataArray:
+            return self.to_xarray()
+        raise DrbChunkError(
+            f"Unsupported chunk impl {impl!r}. Use numpy.ndarray or "
+            f"xarray.DataArray."
+        )
+
+    def _materialise_numpy(self):
+        keys = list(self._manifest.keys())
+        if len(keys) != 1:
+            raise DrbChunkError(
+                f"Direct numpy materialisation needs a single tile; this "
+                f"selection spans {len(keys)} tiles. Use to_xarray() (lazy "
+                f"xarray) for multi-tile assembly."
+            )
+        ref = self._manifest.ref(keys[0])
+        reader = select_reader(ref, self.array, hint=self._reader)
+        resolved = self.array.scheme.resolve(self._selection, self.array) \
+            if self._selection is not None else None
+        return reader.read(ref, resolved)
+
+    def to_xarray(self):
+        """Dask-backed lazy xarray.DataArray assembled from the manifest."""
+        import dask.array as da
+        import xarray as xr
+
+        resolved = (
+            self.array.scheme.resolve(self._selection, self.array)
+            if self._selection is not None else None
+        )
+        keys = list(self._manifest.keys())
+        if not keys:
+            raise DrbChunkError("Empty chunk selection.")
+        blocks = []
+        for key in keys:
+            ref = self._manifest.ref(key)
+            reader = select_reader(ref, self.array, hint=self._reader)
+            shape = tuple(stop - start for (start, stop) in ref.window)
+            delayed = da.from_delayed(
+                _delayed_read(reader, ref, resolved),
+                shape=_block_shape(ref, shape),
+                dtype=self.array.dtype)
+            blocks.append((key, delayed))
+        # v1: single-tile fast path; multi-tile mosaic is future work.
+        if len(blocks) == 1:
+            data = blocks[0][1]
+        else:
+            raise DrbChunkError(
+                "Multi-tile lazy assembly (concat/mosaic) is not yet "
+                "implemented; select a single tile or one window per tile.")
+        return xr.DataArray(data, dims=("band",) + self.array.dims)
+
+    def locator(self) -> dict:
+        return {
+            "source": getattr(self.node, "path", None) and str(self.node.path),
+            "topic": self._topic_uri,
+            "chunk": self.name,
+            "selection": self._selection.to_dict()
+            if self._selection is not None else None,
+        }
+
+    @classmethod
+    def from_locator(cls, data: dict) -> "Chunk":
+        raise DrbChunkError(
+            "locator round-trip (from_locator) is a v1 deferral; "
+            "tracked as a follow-up."
+        )
+
+
+def _block_shape(ref: ChunkRef, window_shape) -> tuple:
+    # rasterio read returns (bands, h, w); v1 assumes a leading band axis.
+    return (1,) + tuple(window_shape)
+
+
+def _delayed_read(reader, ref, resolved):
+    import dask
+    return dask.delayed(reader.read)(ref, resolved)

drb/chunk/core.py

@@ -0,0 +1,123 @@
+# drb/chunk/core.py
+from typing import Any, List, Optional, Tuple, Union
+
+from drb.addons.addon import Addon
+from drb.core.node import DrbNode
+from drb.topics.topic import DrbTopic
+import drb.topics.resolver as resolver
+
+from drb.chunk.chunk import Chunk
+from drb.chunk.descriptor import ChunkDescriptor, retrieve_chunks
+from drb.chunk.exceptions import DrbChunkError
+from drb.chunk.model import ChunkArray, RegularGridManifest
+
+
+class ChunkAddon(Addon):
+    __instance = None
+
+    def __new__(cls, *args, **kwargs):
+        if cls.__instance is None:
+            cls.__instance = super().__new__(cls, *args, **kwargs)
+        return cls.__instance
+
+    @classmethod
+    def identifier(cls) -> str:
+        return "chunk"
+
+    @classmethod
+    def return_type(cls) -> type:
+        return Chunk
+
+    def can_apply(self, source: DrbTopic) -> bool:
+        return bool(retrieve_chunks(source))
+
+    def available_chunks(self, source: DrbTopic) -> List[Tuple[str, dict]]:
+        return [(cd.name, cd.scheme.to_dict())
+                for cd in retrieve_chunks(source).values()]
+
+    def available_collections(self, source: DrbTopic) -> dict:
+        """Map collection name -> list of chunk names declared for it.
+
+        Chunks without a drb:collection are grouped under the key None.
+        """
+        result: dict = {}
+        for cd in retrieve_chunks(source).values():
+            result.setdefault(cd.collection, []).append(cd.name)
+        return result
+
+    def apply(self, node: DrbNode, *, chunk_name: Optional[str] = None,
+              collection: Optional[str] = None,
+              **kwargs) -> Union[Chunk, List[Chunk]]:
+        topic = kwargs.get("topic") or resolver.resolve(node)[0]
+        descriptors = retrieve_chunks(topic)
+        if not descriptors:
+            raise DrbChunkError(
+                f"Topic '{getattr(topic, 'label', topic)}' declares no "
+                f"drb:chunk.")
+        if chunk_name is not None and collection is not None:
+            raise DrbChunkError(
+                "Pass either chunk_name or collection, not both.")
+        if collection is not None:
+            selected = [cd for cd in descriptors.values()
+                        if cd.collection == collection]
+            if not selected:
+                available = sorted({cd.collection
+                                    for cd in descriptors.values()
+                                    if cd.collection is not None})
+                raise DrbChunkError(
+                    f"No collection '{collection}'. "
+                    f"Available: {', '.join(available)}.")
+            return [self._build(node, topic, cd) for cd in selected]
+        if chunk_name is not None:
+            if chunk_name not in descriptors:
+                raise DrbChunkError(
+                    f"No chunk '{chunk_name}' in topic "
+                    f"'{getattr(topic, 'label', topic)}'. "
+                    f"Available: {', '.join(sorted(descriptors))}."
+                )
+            return self._build(node, topic, descriptors[chunk_name])
+        return [self._build(node, topic, cd) for cd in descriptors.values()]
+
+    def _build(self, node: DrbNode, topic: DrbTopic,
+               cd: ChunkDescriptor) -> Chunk:
+        if cd.selection is not None:
+            raise DrbChunkError(
+                f"Chunk '{cd.name}': descriptor-level default "
+                f"drb:selection is a v1 deferral; apply selections "
+                f"explicitly via Chunk.select()."
+            )
+        source = cd.source.extract(node)
+        if source is None:
+            raise DrbChunkError(
+                f"Chunk '{cd.name}': no node matched its source "
+                f"(drb:source) in the product.")
+        source_node = self._resolve_source(node, source)
+        array = ChunkArray(dims=cd.dims,
+                           shape=self._infer_shape(source_node, cd),
+                           dtype=cd.dtype, scheme=cd.scheme)
+        manifest = RegularGridManifest(array, source_node)
+        return Chunk(name=cd.name, array=array, node=source_node,
+                     manifest=manifest, reader=cd.reader,
+                     topic_uri=getattr(topic, "uri", None))
+
+    @staticmethod
+    def _resolve_source(node: DrbNode, source: Any) -> DrbNode:
+        if isinstance(source, DrbNode):
+            return source
+        if source in (".", "", None):
+            return node
+        # A path/URL string -> resolve it against DRB.
+        return resolver.create(str(source))
+
+    @staticmethod
+    def _infer_shape(
+            source_node: DrbNode, cd: ChunkDescriptor) -> Tuple[int, ...]:
+        """Read array shape from the source's existing rasterio impl.
+        No driver change: uses get_impl(DatasetReader).height/width/count."""
+        try:
+            from rasterio.io import DatasetReader
+            ds = source_node.get_impl(DatasetReader)
+            return (ds.height, ds.width)
+        except Exception as exc:        # noqa: BLE001
+            raise DrbChunkError(
+                f"Cannot infer shape for chunk '{cd.name}': {exc}") from exc

drb/chunk/descriptor.py

@@ -0,0 +1,144 @@
+# drb/chunk/descriptor.py
+from dataclasses import dataclass, field
+from typing import Dict, List, Optional, Tuple
+
+import rdflib
+from rdflib.collection import Collection
+
+from drb.extractor import ConstantExtractor, Extractor, parse_extractor
+from drb.exceptions.core import DrbException
+from drb.topics.dao import ManagerDao
+from drb.topics.topic import DrbTopic
+
+from drb.chunk.exceptions import DrbChunkError
+from drb.chunk.selection import Selection, parse_selection
+from drb.chunk.tiling import RegularGrid, TilingScheme
+
+DRB = rdflib.Namespace("http://www.gael.fr/drb#")
+
+
+@dataclass
+class ChunkDescriptor:
+    name: str
+    source: Extractor
+    dims: Tuple[str, ...]
+    dtype: str
+    scheme: TilingScheme
+    selection: Optional[Selection] = None
+    reader: Optional[str] = None
+    collection: Optional[str] = None
+    aux_data: dict = field(default_factory=dict)
+
+
+def _rdf_list(graph: rdflib.Graph, node) -> list:
+    return [str(item) for item in Collection(graph, node)]
+
+
+def _scheme_from_chunk(graph: rdflib.Graph, chunk) -> TilingScheme:
+    scheme_name = graph.value(chunk, DRB.tilingScheme)
+    scheme_name = str(scheme_name) if scheme_name is not None else "regular"
+    if scheme_name != "regular":
+        raise DrbChunkError(
+            f"Unsupported tilingScheme '{scheme_name}'. v1 supports 'regular'."
+        )
+    shape_node = graph.value(chunk, DRB.chunkShape)
+    if shape_node is not None:
+        chunk_shape = tuple(int(v) for v in _rdf_list(graph, shape_node))
+    else:
+        height = graph.value(chunk, DRB.tileHeight)
+        width = graph.value(chunk, DRB.tileWidth)
+        if height is None or width is None:
+            raise DrbChunkError(
+                "RegularGrid needs drb:chunkShape or drb:tileWidth+tileHeight."
+            )
+        chunk_shape = (int(height), int(width))   # dims are (y, x)
+    return RegularGrid(chunk_shape=chunk_shape)
+
+
+# Typed extractor predicates on a drb:source blank node, mapped to the
+# drb.extractor parse_extractor factory keys.
+_SOURCE_KEYS = (
+    ("xquery", DRB.xquery),
+    ("python", DRB.python),
+    ("script", DRB.script),
+    ("constant", DRB.constant),
+)
+
+
+def _source_extractor(graph: rdflib.Graph, chunk) -> Extractor:
+    """Build the source Extractor for one drb:chunk.
+
+    A bare literal stays a ConstantExtractor (v1: "." = the node itself,
+    or a path/URL resolved by the add-on). A typed blank node
+    ``drb:source [ drb:xquery "…" ]`` (or drb:python/script/constant) is
+    dispatched through drb.extractor.parse_extractor; an XQuery navigates
+    the product node and returns the band DrbNode.
+    """
+    value = graph.value(chunk, DRB.source)
+    if value is None:
+        raise DrbChunkError("drb:chunk: missing drb:source.")
+    if isinstance(value, rdflib.Literal):
+        return ConstantExtractor(str(value))
+    for key, predicate in _SOURCE_KEYS:
+        obj = graph.value(value, predicate)
+        if obj is not None:
+            return parse_extractor({key: str(obj)})
+    raise DrbChunkError(
+        "drb:source blank node must carry one of drb:xquery, drb:python, "
+        "drb:script or drb:constant."
+    )
+
+
+def _load_chunks_from_graph(uri: str,
+                            graph: rdflib.Graph) -> List[ChunkDescriptor]:
+    subject = rdflib.URIRef(uri)
+    descriptors: List[ChunkDescriptor] = []
+    for chunk in graph.objects(subject, DRB.chunk):
+        name = graph.value(chunk, DRB.chunkName)
+        dims_node = graph.value(chunk, DRB.dims)
+        dtype = graph.value(chunk, DRB.dtype)
+        reader = graph.value(chunk, DRB.reader)
+        collection_node = graph.value(chunk, DRB.collection)
+        sel_node = graph.value(chunk, DRB.selection)
+        descriptors.append(ChunkDescriptor(
+            name=str(name),
+            source=_source_extractor(graph, chunk),
+            dims=tuple(_rdf_list(graph, dims_node)) if dims_node else (),
+            dtype=str(dtype) if dtype is not None else None,
+            scheme=_scheme_from_chunk(graph, chunk),
+            selection=_selection_from_node(graph, sel_node),
+            reader=str(reader) if reader is not None else None,
+            collection=(str(collection_node)
+                        if collection_node is not None
+                        else None),
+        ))
+    return descriptors
+
+
+def _selection_from_node(graph: rdflib.Graph, node) -> Optional[Selection]:
+    if node is None:
+        return None
+    bands = graph.value(node, DRB.band)
+    if bands is not None:
+        return parse_selection({"band": [int(bands)]})
+    return None
+
+
+def retrieve_chunks(topic: DrbTopic) -> Dict[str, ChunkDescriptor]:
+    """Mirror MetadataAddon._retrieve_metadata: inherit via subClassOf, then
+    override with the topic's own chunks, reading the public dao.graph."""
+    chunks: Dict[str, ChunkDescriptor] = {}
+    if topic.subClassOf is not None:
+        for parent_id in topic.subClassOf:
+            parent = ManagerDao().get_drb_topic(parent_id)
+            chunks.update(retrieve_chunks(parent))
+    if topic.uri is not None:
+        try:
+            dao = ManagerDao().get_dao(topic.id)
+        except DrbException:
+            dao = None
+        graph = getattr(dao, "graph", None)
+        if graph is not None:
+            for cd in _load_chunks_from_graph(topic.uri, graph):
+                chunks[cd.name] = cd
+    return chunks

drb/chunk/exceptions.py

@@ -0,0 +1,9 @@
+from drb.exceptions.core import DrbException
+
+
+class DrbChunkError(DrbException):
+    """Base error for the drb-chunk add-on."""
+
+
+class DrbSelectionError(DrbChunkError):
+    """Raised on an unknown selection type or an out-of-bounds region."""

drb/chunk/interop.py

@@ -0,0 +1,29 @@
+from drb.chunk.chunk import Chunk
+from drb.chunk.exceptions import DrbChunkError
+
+
+def to_kerchunk(chunk: Chunk) -> dict:
+    """Emit the manifest as a kerchunk reference-spec v1 dict.
+
+    kerchunk addresses chunks by byte range in a source file, so only chunks
+    with a ChunkRef.byte_range can be exported. Window-only (format-native)
+    chunks are not representable and raise DrbChunkError.
+    """
+    refs = {}
+    for key in chunk.manifest.keys():
+        ref = chunk.manifest.ref(key)
+        if ref.byte_range is None:
+            raise DrbChunkError(
+                "to_kerchunk requires byte-range chunks; this chunk uses "
+                "format-native windows. Export is not supported."
+            )
+        offset, length = ref.byte_range
+        source = getattr(ref.source, "path", None)
+        if source is None:
+            raise DrbChunkError(
+                f"Chunk source for key {key} has no 'path'; cannot export a "
+                f"kerchunk reference."
+            )
+        key_str = ".".join(str(k) for k in key)
+        refs[key_str] = [str(source), offset, length]
+    return {"version": 1, "refs": refs}

drb/chunk/model.py

@@ -0,0 +1,88 @@
+import abc
+from dataclasses import dataclass
+from typing import Iterable, Optional, Tuple
+
+from drb.core.node import DrbNode
+
+from drb.chunk.exceptions import DrbChunkError
+from drb.chunk.selection import ResolvedSelection
+from drb.chunk.tiling import RegularGrid, TilingScheme
+
+Window = Tuple[Tuple[int, int], ...]
+"""per-dim (start, stop) — geometry-agnostic"""
+
+
+@dataclass
+class ChunkArray:
+    dims: Tuple[str, ...]
+    shape: Tuple[int, ...]
+    dtype: str
+    scheme: TilingScheme
+    coords: Optional[dict] = None
+    attrs: Optional[dict] = None
+
+    @property
+    def grid_shape(self) -> Tuple[int, ...]:
+        if isinstance(self.scheme, RegularGrid):
+            return self.scheme.grid_shape(self)
+        raise DrbChunkError(
+            f"grid_shape is only defined for RegularGrid, not "
+            f"{type(self.scheme).__name__}."
+        )
+
+
+@dataclass
+class ChunkRef:
+    key: Tuple[int, ...]
+    source: DrbNode
+    byte_range: Optional[Tuple[int, int]] = None
+    window: Optional[Window] = None
+    codec: Optional[dict] = None
+
+
+class ChunkManifest(abc.ABC):
+    """Lazy Mapping[key -> ChunkRef] over the scheme's keys."""
+
+    @abc.abstractmethod
+    def ref(self, key: Tuple[int, ...]) -> ChunkRef:
+        ...
+
+    @abc.abstractmethod
+    def keys(self) -> Iterable[Tuple[int, ...]]:
+        ...
+
+    def subset(self, resolved: ResolvedSelection) -> "ChunkManifest":
+        return _SubsetManifest(self, list(resolved.keys))
+
+
+class RegularGridManifest(ChunkManifest):
+    """Manifest over a RegularGrid; windows derived from key * chunk_shape."""
+
+    def __init__(self, array: ChunkArray, source: DrbNode):
+        if not isinstance(array.scheme, RegularGrid):
+            raise DrbChunkError("RegularGridManifest requires a RegularGrid.")
+        self._array = array
+        self._source = source
+        self._chunk_shape = array.scheme.chunk_shape
+
+    def ref(self, key: Tuple[int, ...]) -> ChunkRef:
+        window = tuple(
+            (k * c, min((k + 1) * c, s))
+            for k, c, s in zip(key, self._chunk_shape, self._array.shape)
+        )
+        return ChunkRef(key=key, source=self._source, window=window)
+
+    def keys(self) -> Iterable[Tuple[int, ...]]:
+        return self._array.scheme.keys(self._array)
+
+
+class _SubsetManifest(ChunkManifest):
+    def __init__(self, parent: ChunkManifest, keys):
+        self._parent = parent
+        self._keys = keys
+
+    def ref(self, key: Tuple[int, ...]) -> ChunkRef:
+        return self._parent.ref(key)
+
+    def keys(self) -> Iterable[Tuple[int, ...]]:
+        return iter(self._keys)

drb/chunk/readers.py

@@ -0,0 +1,70 @@
+# drb/chunk/readers.py
+import abc
+import io
+from typing import Any, Optional
+
+from drb.chunk.exceptions import DrbChunkError
+from drb.chunk.model import ChunkArray, ChunkRef
+from drb.chunk.selection import ResolvedSelection
+
+
+class ReaderStrategy(abc.ABC):
+    """Materialises one chunk using source node's existing get_impl."""
+
+    @abc.abstractmethod
+    def can_read(self, ref: ChunkRef, array: ChunkArray) -> bool:
+        ...
+
+    @abc.abstractmethod
+    def read(self, ref: ChunkRef, resolved: ResolvedSelection) -> Any:
+        ...
+
+
+class RasterWindowReader(ReaderStrategy):
+    """Windowed rasterio read via the source's EXISTING get_impl. The driver
+    (impl/image) is not modified."""
+
+    def can_read(self, ref: ChunkRef, array: ChunkArray) -> bool:
+        return ref.window is not None
+
+    def read(self, ref: ChunkRef, resolved: ResolvedSelection) -> Any:
+        from rasterio.io import DatasetReader
+        dataset = ref.source.get_impl(DatasetReader)
+        # Pass window as tuple-of-ranges; rasterio normalises it internally.
+        return dataset.read(window=ref.window)
+
+
+class ByteRangeReader(ReaderStrategy):
+    """Reads ref.byte_range via the source driver's partial GET."""
+
+    def can_read(self, ref: ChunkRef, array: ChunkArray) -> bool:
+        return ref.byte_range is not None
+
+    def read(self, ref: ChunkRef, resolved: ResolvedSelection) -> Any:
+        offset, length = ref.byte_range
+        stream = ref.source.get_impl(io.BytesIO)
+        stream.seek(offset)
+        return stream.read(length)
+
+
+_STRATEGIES = (RasterWindowReader(), ByteRangeReader())
+_BY_NAME = {"raster": RasterWindowReader, "range": ByteRangeReader}
+
+
+def select_reader(ref: ChunkRef, array: ChunkArray,
+                  hint: Optional[str] = None) -> ReaderStrategy:
+    if hint is not None:
+        try:
+            return _BY_NAME[hint]()
+        except KeyError as exc:
+            raise DrbChunkError(
+                f"Unknown reader hint '{hint}'. "
+                f"Available: {', '.join(sorted(_BY_NAME))}."
+            ) from exc
+    for strategy in _STRATEGIES:
+        if strategy.can_read(ref, array):
+            return strategy
+    raise DrbChunkError(
+        "No reader strategy applies to this chunk (no window, no byte_range). "
+        "Refusing to fall back to a full read."
+    )