Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
3 changes: 3 additions & 0 deletions CHANGES.md
Original file line number Diff line number Diff line change
Expand Up @@ -2,6 +2,9 @@

## Unreleased

### titiler.xarray
* add `opener_options` arg to `titiler.xarray.io.Reader` to allow users to pass args through to a custom opener function ([#1248(https://github.com/developmentseed/titiler/pull/1248)])

### Misc

* update docker image to python:3.13
Expand Down
65 changes: 65 additions & 0 deletions src/titiler/xarray/tests/test_io_tools.py
Original file line number Diff line number Diff line change
Expand Up @@ -2,6 +2,8 @@

import os
from datetime import datetime
from typing import Any, Dict, Optional
from urllib.parse import urlparse

import numpy
import pytest
Expand Down Expand Up @@ -223,6 +225,69 @@ def test_reader(protocol, filename):
assert src.tile(0, 0, 0)


def test_opener():
"""test custom opener"""
src_path = "file://" + os.path.join("file://", prefix, "dataset_2d.nc")

def custom_netcdf_opener( # noqa: C901
src_path: str,
special_arg: bool,
group: Optional[str] = None,
decode_times: bool = True,
) -> xarray.Dataset:
"""Open Xarray dataset with fsspec.

Args:
src_path (str): dataset path.
group (Optional, str): path to the netCDF/Zarr group in the given file to open given as a str.
decode_times (bool): If True, decode times encoded in the standard NetCDF datetime format into datetime objects. Otherwise, leave them encoded as numbers.

Returns:
xarray.Dataset

"""
import fsspec # noqa

parsed = urlparse(src_path)
protocol = parsed.scheme or "file"

if not special_arg:
raise ValueError("you forgot the special_arg :(")

xr_open_args: Dict[str, Any] = {
"decode_coords": "all",
"decode_times": decode_times,
"engine": "h5netcdf",
"lock": False,
}

# Argument if we're opening a datatree
if group is not None:
xr_open_args["group"] = group

fs = fsspec.filesystem(protocol)
ds = xarray.open_dataset(fs.open(src_path), **xr_open_args)

return ds

with Reader(
src_path=src_path,
opener=custom_netcdf_opener,
opener_options={"special_arg": True},
variable="dataset",
) as src:
assert src.info()

with pytest.raises(ValueError):
with Reader(
src_path=src_path,
opener=custom_netcdf_opener,
opener_options={"special_arg": False},
variable="dataset",
) as src:
pass


@pytest.mark.parametrize(
"group",
[0, 1, 2],
Expand Down
2 changes: 2 additions & 0 deletions src/titiler/xarray/titiler/xarray/io.py
Original file line number Diff line number Diff line change
Expand Up @@ -196,6 +196,7 @@ class Reader(XarrayReader):

# xarray.Dataset options
opener: Callable[..., xarray.Dataset] = attr.ib(default=xarray_open_dataset)
opener_options: Dict = attr.ib(factory=dict)

group: Optional[str] = attr.ib(default=None)
decode_times: bool = attr.ib(default=True)
Expand All @@ -219,6 +220,7 @@ def __attrs_post_init__(self):
self.src_path,
group=self.group,
decode_times=self.decode_times,
**self.opener_options,
)

self.input = get_variable(
Expand Down
Loading