Extendable entrypoint plugins

.pre-commit-config.yaml

@@ -10,7 +10,7 @@ repos:
       - id: double-quote-string-fixer
 
   - repo: https://github.com/psf/black
-    rev: 22.10.0
+    rev: 22.12.0
     hooks:
       - id: black
         args: ["--line-length", "100", "--skip-string-normalization"]

setup.cfg

@@ -9,7 +9,7 @@ select = B,C,E,F,W,T4,B9
 
 [isort]
 known_first_party=xpublish
-known_third_party=cachey,dask,fastapi,numcodecs,numpy,pandas,pkg_resources,pytest,setuptools,sphinx_autosummary_accessors,starlette,uvicorn,xarray,zarr
+known_third_party=cachey,dask,fastapi,numcodecs,numpy,pandas,pkg_resources,pydantic,pytest,setuptools,sphinx_autosummary_accessors,starlette,uvicorn,xarray,zarr
 multi_line_output=3
 include_trailing_comma=True
 force_grid_wrap=0

setup.py

@@ -46,4 +46,12 @@
     keywords=['xarray', 'zarr', 'api'],
     use_scm_version={'version_scheme': 'post-release', 'local_scheme': 'dirty-tag'},
     setup_requires=['setuptools_scm>=3.4', 'setuptools>=42'],
+    entry_points={
+        'xpublish.plugin': [
+            'info = xpublish.included_plugins.dataset_info:DatasetInfoPlugin',
+            'zarr = xpublish.included_plugins.zarr:ZarrPlugin',
+            'module_version = xpublish.included_plugins.module_version:ModuleVersionPlugin',
+            'plugin_info = xpublish.included_plugins.plugin_info:PluginInfoPlugin',
+        ]
+    },
 )

tests/test_rest_api.py

@@ -102,7 +102,7 @@ def test_custom_app_routers(airtemp_ds, dims_router, router_kws, path):
     else:
         routers = [(dims_router, router_kws)]
 
-    rest = Rest(airtemp_ds, routers=routers)
+    rest = Rest(airtemp_ds, routers=routers, plugins={})
     client = TestClient(rest.app)
 
     response = client.get(path)

xpublish/__init__.py

@@ -1,6 +1,8 @@
 from pkg_resources import DistributionNotFound, get_distribution
 
-from .rest import Rest, RestAccessor  # noqa: F401
+from .accessor import RestAccessor  # noqa: F401
+from .plugin import Plugin, Router  # noqa: F401
+from .rest import Rest  # noqa: F401
 
 try:
     __version__ = get_distribution(__name__).version

xpublish/accessor.py

@@ -0,0 +1,74 @@
+import cachey
+import xarray as xr
+from fastapi import FastAPI
+
+from .rest import Rest
+
+
+@xr.register_dataset_accessor('rest')
+class RestAccessor:
+    """REST API Accessor for serving one dataset in its
+    dedicated FastAPI application.
+
+    """
+
+    def __init__(self, xarray_obj):
+
+        self._obj = xarray_obj
+        self._rest = None
+
+        self._initialized = False
+
+    def _get_rest_obj(self):
+        if self._rest is None:
+            self._rest = Rest(self._obj)
+
+        return self._rest
+
+    def __call__(self, **kwargs):
+        """Initialize this accessor by setting optional configuration values.
+
+        Parameters
+        ----------
+        **kwargs
+            Arguments passed to :func:`xpublish.Rest.__init__`.
+
+        Notes
+        -----
+        This method can only be invoked once.
+
+        """
+        if self._initialized:
+            raise RuntimeError('This accessor has already been initialized')
+        self._initialized = True
+
+        self._rest = Rest(self._obj, **kwargs)
+
+        return self
+
+    @property
+    def cache(self) -> cachey.Cache:
+        """Returns the :class:`cachey.Cache` instance used by the FastAPI application."""
+
+        return self._get_rest_obj().cache
+
+    @property
+    def app(self) -> FastAPI:
+        """Returns the :class:`fastapi.FastAPI` application instance."""
+
+        return self._get_rest_obj().app
+
+    def serve(self, **kwargs):
+        """Serve this FastAPI application via :func:`uvicorn.run`.
+
+        Parameters
+        ----------
+        **kwargs :
+            Arguments passed to :func:`xpublish.Rest.serve`.
+
+        Notes
+        -----
+        This method is blocking and does not return.
+
+        """
+        self._get_rest_obj().serve(**kwargs)

xpublish/dependencies.py

@@ -1,15 +1,20 @@
 """
 Helper functions to use a FastAPI dependencies.
 """
+from typing import TYPE_CHECKING, Dict, List
+
 import cachey
 import xarray as xr
 from fastapi import Depends
 
 from .utils.api import DATASET_ID_ATTR_KEY
 from .utils.zarr import create_zmetadata, create_zvariables, zarr_metadata_key
 
+if TYPE_CHECKING:
+    from .plugin import Plugin
+
 
-def get_dataset_ids():
+def get_dataset_ids() -> List[str]:
     """FastAPI dependency for getting the list of ids (string keys)
     of the collection of datasets being served.
 
@@ -23,7 +28,7 @@ def get_dataset_ids():
     return []  # pragma: no cover
 
 
-def get_dataset(dataset_id: str):
+def get_dataset(dataset_id: str) -> xr.Dataset:
     """FastAPI dependency for accessing the published xarray dataset object.
 
     Use this callable as dependency in any FastAPI path operation
@@ -33,10 +38,10 @@ def get_dataset(dataset_id: str):
     application.
 
     """
-    return None  # pragma: no cover
+    return xr.Dataset()
 
 
-def get_cache():
+def get_cache() -> cachey.Cache:
     """FastAPI dependency for accessing the application's cache.
 
     Use this callable as dependency in any FastAPI path operation
@@ -47,7 +52,7 @@ def get_cache():
     application.
 
     """
-    return None  # pragma: no cover
+    return cachey.Cache(available_bytes=1e6)
 
 
 def get_zvariables(
@@ -84,3 +89,9 @@ def get_zmetadata(
         cache.put(cache_key, zmeta, 99999)
 
     return zmeta
+
+
+def get_plugins() -> Dict[str, 'Plugin']:
+    """FastAPI dependency that returns the a dictionary of loaded plugins"""
+
+    return []

xpublish/included_plugins/dataset_info.py

@@ -0,0 +1,71 @@
+import xarray as xr
+from fastapi import Depends
+from pydantic import Field
+from starlette.responses import HTMLResponse
+from zarr.storage import attrs_key
+
+from ..dependencies import get_zmetadata, get_zvariables
+from ..plugin import Plugin, Router
+
+
+class DatasetInfoRouter(Router):
+    """API entry-points providing basic information about the dataset(s)."""
+
+    prefix = ''
+
+    def register(self):
+        @self._router.get('/')
+        def html_representation(
+            dataset=Depends(self.deps.dataset),
+        ):
+            """Returns a HTML representation of the dataset."""
+
+            with xr.set_options(display_style='html'):
+                return HTMLResponse(dataset._repr_html_())
+
+        @self._router.get('/keys')
+        def list_keys(
+            dataset=Depends(self.deps.dataset),
+        ):
+            return list(dataset.variables)
+
+        @self._router.get('/dict')
+        def to_dict(
+            dataset=Depends(self.deps.dataset),
+        ):
+            return dataset.to_dict(data=False)
+
+        @self._router.get('/info')
+        def info(
+            dataset=Depends(self.deps.dataset),
+            cache=Depends(self.deps.cache),
+        ):
+            """Dataset schema (close to the NCO-JSON schema)."""
+
+            zvariables = get_zvariables(dataset, cache)
+            zmetadata = get_zmetadata(dataset, cache, zvariables)
+
+            info = {}
+            info['dimensions'] = dict(dataset.dims.items())
+            info['variables'] = {}
+
+            meta = zmetadata['metadata']
+
+            for name, var in zvariables.items():
+                attrs = meta[f'{name}/{attrs_key}']
+                attrs.pop('_ARRAY_DIMENSIONS')
+                info['variables'][name] = {
+                    'type': var.data.dtype.name,
+                    'dimensions': list(var.dims),
+                    'attributes': attrs,
+                }
+
+            info['global_attributes'] = meta[attrs_key]
+
+            return info
+
+
+class DatasetInfoPlugin(Plugin):
+    name = 'dataset_info'
+
+    dataset_router: DatasetInfoRouter = Field(default_factory=DatasetInfoRouter)

xpublish/included_plugins/module_version.py

@@ -0,0 +1,49 @@
+"""
+Version information router
+"""
+import importlib
+import sys
+
+from pydantic import Field
+
+from ..plugin import Plugin, Router
+from ..utils.info import get_sys_info, netcdf_and_hdf5_versions
+
+
+class ModuleVersionAppRouter(Router):
+    """Module and system version information"""
+
+    prefix = ''
+
+    def register(self):
+        @self._router.get('/versions')
+        def get_versions():
+            versions = dict(get_sys_info() + netcdf_and_hdf5_versions())
+            modules = [
+                'xarray',
+                'zarr',
+                'numcodecs',
+                'fastapi',
+                'starlette',
+                'pandas',
+                'numpy',
+                'dask',
+                'distributed',
+                'uvicorn',
+            ]
+            for modname in modules:
+                try:
+                    if modname in sys.modules:
+                        mod = sys.modules[modname]
+                    else:  # pragma: no cover
+                        mod = importlib.import_module(modname)
+                    versions[modname] = getattr(mod, '__version__', None)
+                except ImportError:  # pragma: no cover
+                    pass
+            return versions
+
+
+class ModuleVersionPlugin(Plugin):
+    name = 'module_version'
+
+    app_router: ModuleVersionAppRouter = Field(default_factory=ModuleVersionAppRouter)

xpublish/included_plugins/plugin_info.py

@@ -0,0 +1,49 @@
+"""
+Plugin information router
+"""
+import importlib
+from typing import Dict, Optional
+
+from fastapi import Depends
+from pydantic import BaseModel, Field
+
+from ..plugin import Plugin, Router
+
+
+class PluginInfo(BaseModel):
+    path: str
+    version: Optional[str]
+
+
+class PluginInfoAppRouter(Router):
+    """Plugin information"""
+
+    prefix = ''
+
+    def register(self):
+        @self._router.get('/plugins')
+        def get_plugins(
+            plugins: Dict[str, Plugin] = Depends(self.deps.plugins)
+        ) -> Dict[str, PluginInfo]:
+            plugin_info = {}
+
+            for name, plugin in plugins.items():
+                plugin_type = type(plugin)
+                module_name = plugin_type.__module__.split('.')[0]
+                try:
+                    mod = importlib.import_module(module_name)
+                    version = getattr(mod, '__version__', None)
+                except ImportError:
+                    version = None
+
+                plugin_info[name] = PluginInfo(
+                    path=f'{plugin_type.__module__}.{plugin.__repr_name__()}', version=version
+                )
+
+            return plugin_info
+
+
+class PluginInfoPlugin(Plugin):
+    name = 'plugin_info'
+
+    app_router: PluginInfoAppRouter = Field(default_factory=PluginInfoAppRouter)