Skip to content

Instantly share code, notes, and snippets.

@benbovy
Last active June 17, 2020 15:24
Show Gist options
  • Save benbovy/f687e1697e18e39e4924872cd83a60e9 to your computer and use it in GitHub Desktop.
Save benbovy/f687e1697e18e39e4924872cd83a60e9 to your computer and use it in GitHub Desktop.
FastAPI flexible routes
from typing import Callable, Dict, Union
from fastapi import APIRouter, Depends, FastAPI, HTTPException
import xarray as xr
DatasetOrCollection = Union[xr.Dataset, Dict[str, xr.Dataset]]
def _get_dataset_dependency(obj: DatasetOrCollection) -> Callable:
"""Returns a xarray Dataset getter to be used as fastAPI dependency."""
def get_obj():
return obj
def get_from_mapping(dataset_id: str):
if dataset_id not in obj:
raise HTTPException(status_code=404, detail="Dataset not found")
return obj[dataset_id]
if isinstance(obj, xr.Dataset):
return get_obj
else:
return get_from_mapping
class APIRouterWrapper:
"""Wraps :class:`fastapi.APIRouter` so that it can be included
in an application serving either a single xarray Dataset or a
collection of Datasets.
"""
def __init__(self, obj: DatasetOrCollection):
self._obj = obj
self._router = None
self._dataset = None
@property
def dataset(self) -> Callable:
if self._dataset is None:
self._dataset = _get_dataset_dependency(self._obj)
return self._dataset
def init_router(self):
self._router = APIRouter()
@property
def router(self) -> APIRouter:
if self._router is None:
self.init_router()
return self._router
class ExampleRouter(APIRouterWrapper):
"""A simple example."""
def init_router(self):
super().init_router()
@self._router.get("/dims")
def get_dims(dataset: xr.Dataset = Depends(self.dataset)):
return dataset.dims
# -- reuse ExampleRouter in an app serving one dataset
single_app = FastAPI()
r1 = ExampleRouter(xr.Dataset({'x': [1, 2, 3]}))
single_app.include_router(r1.router, prefix='')
# -- reuse ExampleRouter in an app serving a collection of datasets
multi_app = FastAPI()
r2 = ExampleRouter({
'd1': xr.Dataset({'x': [1, 2, 3]}),
'd2': xr.Dataset({'x': [1, 2, 3], 'y': [4, 5, 6]})
})
multi_app.include_router(r2.router, prefix='/datasets/{dataset_id}')
# -- main app with both examples above mounted
app = FastAPI()
app.mount("/single", single_app)
app.mount("/multi", multi_app)
from typing import Optional
from fastapi import APIRouter, Depends, FastAPI, HTTPException
import xarray as xr
def get_dataset(dataset_id: Optional[str] = None):
"""FastAPI dependency for accessing a xarray dataset object.
Use this callable as dependency in any FastAPI path operation
function where you need access to the xarray Dataset being served.
This dummy dependency will be overridden when creating the FastAPI
application.
"""
return None
router = APIRouter()
@router.get("/dims")
def get_dims(dataset: xr.Dataset = Depends(get_dataset)):
return dataset.dims
# -- reuse router in an app serving one dataset
ds = xr.Dataset({'x': [1, 2, 3]})
single_app = FastAPI()
single_app.include_router(router, prefix='')
single_app.dependency_overrides[get_dataset] = lambda : ds
# -- reuse router in an app serving a collection of datasets
ds_collection = {
'd1': xr.Dataset({'x': [1, 2, 3]}),
'd2': xr.Dataset({'x': [1, 2, 3], 'y': [4, 5, 6]})
}
def get_dataset_from_dict(dataset_id: str):
if dataset_id not in ds_collection:
raise HTTPException(status_code=404, detail="Dataset not found")
return ds_collection[dataset_id]
multi_app = FastAPI()
multi_app.include_router(router, prefix='/datasets/{dataset_id}')
multi_app.dependency_overrides[get_dataset] = get_dataset_from_dict
# -- main app with both examples above mounted
app = FastAPI()
app.mount("/single", single_app)
app.mount("/multi", multi_app)
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment