Source code for tiled.adapters.xarray

import itertools
from collections.abc import Mapping
from typing import Any, Iterator, List, Optional

import xarray

from ..structures.core import Spec
from .array import ArrayAdapter
from .mapping import MapAdapter


class DatasetAdapter(MapAdapter):
    """
    Wrap an xarray.Dataset
    """

[docs] @classmethod def from_dataset( cls, dataset: Any, *, specs: Optional[List[Spec]] = None, ) -> "DatasetAdapter": """ Parameters ---------- dataset : specs : Returns ------- """ mapping = _DatasetMap(dataset) specs = specs or [] if "xarray_dataset" not in [spec.name for spec in specs]: specs.append(Spec("xarray_dataset")) return cls( mapping, metadata={"attrs": dataset.attrs}, specs=specs, )
def __init__( self, mapping: Any, *args: Any, specs: Optional[List[Spec]] = None, **kwargs: Any, ) -> None: """ Parameters ---------- mapping : args : specs : kwargs : """ if isinstance(mapping, xarray.Dataset): raise TypeError( "Use DatasetAdapter.from_dataset(...), not DatasetAdapter(...)." ) super().__init__(mapping, *args, specs=specs, **kwargs) def inlined_contents_enabled(self, depth: int) -> bool: """ Parameters ---------- depth : Returns ------- """ # Tell the server to in-line the description of each array # (i.e. data_vars and coords) to avoid latency of a second # request. return True class _DatasetMap(Mapping[str, Any]): def __init__(self, dataset: Any) -> None: """ Parameters ---------- dataset : """ self._dataset = dataset def __len__(self) -> int: """ Returns ------- """ return len(self._dataset.data_vars) + len(self._dataset.coords) def __iter__(self) -> Iterator[Any]: """ Returns ------- """ yield from itertools.chain(self._dataset.data_vars, self._dataset.coords) def __getitem__(self, key: str) -> ArrayAdapter: """ Parameters ---------- key : Returns ------- """ data_array = self._dataset[key] if key in self._dataset.coords: spec = Spec("xarray_coord") else: spec = Spec("xarray_data_var") return ArrayAdapter.from_array( data_array.data, metadata={"attrs": data_array.attrs}, dims=data_array.dims, specs=[spec], )