FastAPI flexible routes

flexible_routes.py

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)

flexible_routes_alt.py

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)