Extend xr.open_dataset

Author: jsignell

.gitignore

@@ -1,129 +1,5 @@
-# Byte-compiled / optimized / DLL files
-__pycache__/
-*.py[cod]
-*$py.class
-
-# C extensions
-*.so
-
-# Distribution / packaging
-.Python
-build/
-develop-eggs/
-dist/
-downloads/
-eggs/
-.eggs/
-lib/
-lib64/
-parts/
-sdist/
-var/
-wheels/
-pip-wheel-metadata/
-share/python-wheels/
-*.egg-info/
-.installed.cfg
-*.egg
-MANIFEST
-
-# PyInstaller
-#  Usually these files are written by a python script from a template
-#  before PyInstaller builds the exe, so as to inject date/other infos into it.
-*.manifest
-*.spec
-
-# Installer logs
-pip-log.txt
-pip-delete-this-directory.txt
-
-# Unit test / coverage reports
-htmlcov/
-.tox/
-.nox/
+__pycache__
 .coverage
-.coverage.*
-.cache
-nosetests.xml
-coverage.xml
-*.cover
-*.py,cover
-.hypothesis/
-.pytest_cache/
-
-# Translations
-*.mo
-*.pot
-
-# Django stuff:
-*.log
-local_settings.py
-db.sqlite3
-db.sqlite3-journal
-
-# Flask stuff:
-instance/
-.webassets-cache
-
-# Scrapy stuff:
-.scrapy
-
-# Sphinx documentation
-docs/_build/
-
-# PyBuilder
-target/
-
-# Jupyter Notebook
-.ipynb_checkpoints
-
-# IPython
-profile_default/
-ipython_config.py
-
-# pyenv
-.python-version
-
-# pipenv
-#   According to pypa/pipenv#598, it is recommended to include Pipfile.lock in version control.
-#   However, in case of collaboration, if having platform-specific dependencies or dependencies
-#   having no cross-platform support, pipenv may install dependencies that don't work, or not
-#   install all needed dependencies.
-#Pipfile.lock
-
-# PEP 582; used by e.g. github.com/David-OConnor/pyflow
-__pypackages__/
-
-# Celery stuff
-celerybeat-schedule
-celerybeat.pid
-
-# SageMath parsed files
-*.sage.py
-
-# Environments
-.env
-.venv
-env/
-venv/
-ENV/
-env.bak/
-venv.bak/
-
-# Spyder project settings
-.spyderproject
-.spyproject
-
-# Rope project settings
-.ropeproject
-
-# mkdocs documentation
-/site
-
-# mypy
-.mypy_cache/
-.dmypy.json
-dmypy.json
-
-# Pyre type checker
-.pyre/
+htmlcov
+dist
+.direnv

.pre-commit-config.yaml

@@ -0,0 +1,12 @@
+
+repos:
+-   repo: https://github.com/psf/black
+    rev: "22.12.0"
+    hooks:
+    -   id: black
+        additional_dependencies: ['click==8.0.4']
+-   repo: https://github.com/pycqa/flake8
+    rev: 6.0.0
+    hooks:
+    - id: flake8
+      language_version: python3

README.md

@@ -1,2 +1,30 @@
-# xarray-pystac-backend
-For extending xarray.open_dataset to accept pystac objects
+# xpystac
+For extending `xarray.open_dataset` to accept pystac objects
+
+## Example
+
+```python
+from pystac_client import Client
+import xarray as xr
+
+
+client = Client.open("https://earth-search.aws.element84.com/v0")
+
+search = client.search(
+    intersects=dict(type="Point", coordinates=[-105.78, 35.79]),
+    collections=['sentinel-s2-l2a-cogs'],
+    datetime="2020-04-01/2020-05-01",
+)
+
+xr.open_dataset(search, engine="stac")
+```
+
+## Install
+
+```bash
+pip install git+https://github.com/jsignell/xpystac
+```
+
+## Prior Art
+
+This work is inspired by https://github.com/TomAugspurger/staccontainers

pyproject.toml

@@ -0,0 +1,27 @@
+[build-system]
+requires = ["flit_core >=2,<4"]
+build-backend = "flit_core.buildapi"
+
+[project]
+name = "xpystac"
+author = "Julia Signell"
+author-email = "[email protected]"
+classifiers = [ "License :: OSI Approved :: MIT License",]
+dependencies = [
+    "xarray",
+    "pystac>=1.0.0b3",
+]
+requires-python = ">=3.8"
+description = "Extend xarray.open_dataset to accept pystac objects"
+license = {text = "MIT"}
+readme = "README.md"
+dynamic = ["version"]
+
+[project.optional-dependencies]
+dev = [
+  "pytest",
+  "pre-commit",
+]
+
+[project.entry-points."xarray.backends"]
+stac = "xpystac.xarray_plugin:STACBackend"

setup.cfg

@@ -0,0 +1,19 @@
+[flake8]
+# References:
+# https://flake8.readthedocs.io/en/latest/user/configuration.html
+# https://flake8.readthedocs.io/en/latest/user/error-codes.html
+# https://pycodestyle.pycqa.org/en/latest/intro.html#error-codes
+exclude = __init__.py
+ignore =
+    E20,
+    E231,
+    E241,
+    E26,
+    E4,
+    E721,
+    E731,
+    E741,
+    W503,
+    W504,
+    F811,
+max-line-length = 120

xpystac/__init__.py

@@ -0,0 +1 @@
+__version__ = "0.0.1"

xpystac/utils.py

@@ -0,0 +1,9 @@
+import importlib
+
+
+def _import_optional_dependency(name):
+    try:
+        module = importlib.import_module(name)
+    except ImportError as e:
+        raise ImportError(f"Missing optional dependency '{name}'") from e
+    return module

xpystac/xarray_plugin.py

@@ -0,0 +1,78 @@
+import functools
+
+import pystac
+import xarray
+from xarray.backends import BackendEntrypoint
+
+from .utils import _import_optional_dependency
+
+
+@functools.singledispatch
+def to_xarray(item, **kwargs) -> xarray.Dataset:
+    if hasattr(item, "get_all_items"):
+        item_collection = item.get_all_items()
+        return to_xarray(item_collection)
+    raise TypeError
+
+
+@to_xarray.register(pystac.Item)
+@to_xarray.register(pystac.ItemCollection)
+def _(
+    obj: pystac.Item | pystac.ItemCollection,
+    drop_variables: str | list[str] = None,
+    **kwargs,
+) -> xarray.Dataset:
+    stackstac = _import_optional_dependency("stackstac")
+    if drop_variables is not None:
+        raise KeyError("``drop_variables`` not implemented for pystac items")
+    return stackstac.stack(obj, **kwargs).to_dataset(dim="band", promote_attrs=True)
+
+
+@to_xarray.register
+def _(obj: pystac.Asset, **kwargs) -> xarray.Dataset:
+    open_kwargs = obj.extra_fields.get("xarray:open_kwargs", {})
+
+    if obj.media_type == pystac.MediaType.JSON and "index" in obj.roles:
+        requests = _import_optional_dependency("requests")
+        fsspec = _import_optional_dependency("fsspec")
+        r = requests.get(obj.href)
+        r.raise_for_status()
+        try:
+            import planetary_computer
+
+            refs = planetary_computer.sign(r.json())
+        except ImportError:
+            refs = r.json()
+        mapper = fsspec.get_mapper("reference://", fo=refs)
+        default_kwargs = {"engine": "zarr", "consolidated": False, "chunks": {}}
+        return xarray.open_dataset(mapper, **default_kwargs, **open_kwargs, **kwargs)
+
+    if obj.media_type == pystac.MediaType.COG:
+        default_kwargs = {"engine", "rasterio"}
+    elif obj.media_type == "application/vnd+zarr":
+        default_kwargs = {"engine": "zarr"}
+    else:
+        default_kwargs = {}
+
+    ds = xarray.open_dataset(obj.href, **default_kwargs, **open_kwargs, **kwargs)
+    return ds
+
+
+class STACBackend(BackendEntrypoint):
+    def open_dataset(
+        self,
+        obj,
+        *,
+        drop_variables: str | list[str] = None,
+        **kwargs,
+    ):
+        return to_xarray(obj, drop_variables=drop_variables, **kwargs)
+
+    open_dataset_parameters = ["obj", "drop_variables"]
+
+    def guess_can_open(self, obj):
+        return isinstance(obj, (pystac.Asset, pystac.Item, pystac.ItemCollection))
+
+    description = "Open pystac objects in Xarray"
+
+    url = "https://github.com/jsignell/xpystac"