Dispatch Wrapper

earthkit-utils provides a dispatch wrapper function which can be used to dispatch a function to the appropriate method given the inputs provided. For example, it is useful to have a FieldList and Xarray version of a high-level function such that the appropriate object preparation can be done prior to executing the array-api function.

How The Wrapper Resolves Implementations

The dispatch wrapper routes calls to sibling modules based on the input type:

• .../api.py contains the high-level function (demonstration module name)

• .../array.py contains the array implementation

• .../xarray.py contains the xarray implementation

Inside the high-level function, the execution can be dispatched to the appropriate object specific function:

dispatched = dispatch(my_function, fieldlist=False, array=True)
return dispatched(arg1, arg2, ...)

Below, we build a minimal runnable demo package in-memory to show this behaviour end-to-end.

import numpy as np
import xarray as xr
from earthkit.data import SimpleFieldList

# Test data
TEST_NUMPY_ARRAY = np.array([1, 2, 3, 4, 5])
TEST_XARRAY_DATAARRAY = xr.DataArray(TEST_NUMPY_ARRAY, name="test", dims=["x"], coords={"x": [0, 1, 2, 3, 4]})
TEST_XARRAY_DATASET = xr.Dataset({"test": TEST_XARRAY_DATAARRAY})
TEST_FIELDLIST = SimpleFieldList(TEST_NUMPY_ARRAY)

Below we define some dummy modules that we can use to demonstrate the wrapper functionality

import sys
import types


def build_demo_modules():
    # Create in-memory modules that mimic a package layout:
    # dispatcher_demo.api, dispatcher_demo.array, dispatcher_demo.xarray
    pkg = types.ModuleType("dispatcher_demo")
    pkg.__path__ = []

    api_mod = types.ModuleType("dispatcher_demo.api")
    array_mod = types.ModuleType("dispatcher_demo.array")
    xarray_mod = types.ModuleType("dispatcher_demo.xarray")
    fieldlist_mod = types.ModuleType("dispatcher_demo.fieldlist")

    # Register package and submodules so import_module() can resolve them
    sys.modules["dispatcher_demo"] = pkg
    sys.modules["dispatcher_demo.api"] = api_mod
    sys.modules["dispatcher_demo.array"] = array_mod
    sys.modules["dispatcher_demo.xarray"] = xarray_mod
    sys.modules["dispatcher_demo.fieldlist"] = fieldlist_mod

    # Low-level implementations
    def add_one_array(t):
        print("Called add_one_array")
        return np.asarray(t) + 1

    def add_one_xarray(t):
        print("Called add_one_xarray")
        return t + 1

    def offset_array(t, bias=0):
        print("Called offset_array")
        return np.asarray(t) + bias

    def offset_xarray(t, bias=0):
        print("Called offset_xarray")
        return t + bias

    array_mod.add_one = add_one_array
    xarray_mod.add_one = add_one_xarray
    array_mod.offset = offset_array
    xarray_mod.offset = offset_xarray

    # High-level API functions using the dispatch wrapper
    api_code = """
from earthkit.utils.decorators import dispatch

def add_one(t):
    dispatched = dispatch(add_one, fieldlist=False, array=True)
    return dispatched(t)

def offset(t, bias=0):
    dispatched = dispatch(offset, match="t", fieldlist=False, array=True)
    return dispatched(t, bias=bias)

def normalize_like(t):
    dispatched = dispatch(normalize_like, fieldlist=False, array=False, array_like=True)
    return dispatched(t)
"""
    exec(api_code, api_mod.__dict__)

    # Add array-like implementation for normalize_like in .array
    def normalize_like_array(t):
        a = np.asarray(t, dtype=float)
        return (a - a.mean()) / (a.std() + 1e-12)

    array_mod.normalize_like = normalize_like_array

    return api_mod


api = build_demo_modules()

We now define some functions which implement the dispatch wrapper.

The add_one function can be dispatched to an array function and an xarray function, but not a fieldlist or array_like function.

The offset function can be dispatched to an array function, an xarray and an array_like function, but not a fieldlist function.

from earthkit.utils.decorators import dispatch


# Define a function that can dispatch to an array function and an xarray function,
# but not a fieldlist or array_like function
def add_one(data):
    dispatched = dispatch(api.add_one, fieldlist=False)  # Default values: xarray=True, array=True, array_like=False
    return dispatched(data)


# Define a function that can dispatch to an array function, an xarray and an array_like function,
# but not a fieldlist function
def offset(data):
    dispatched = dispatch(api.offset, xarray=False, fieldlist=False)  # array=True, array_like=False
    return dispatched(data)


# Define a function that can dispatch to an array function, an xarray and an array_like  function,
# but not a fieldlist function
def offset_array_like(data):
    dispatched = dispatch(api.offset, xarray=False, fieldlist=False, array_like=True)  # array=True
    return dispatched(data)

Below we try to call the functions with a range of data objects, those which match the object types are executed with the appropriate function. Those which do not match, raise an appropriate error message.

for data_object in [TEST_NUMPY_ARRAY, TEST_XARRAY_DATAARRAY, TEST_XARRAY_DATASET, TEST_FIELDLIST, [1, 2, 3]]:
    try:
        result = add_one(data_object)
        print(f"Successful execution for input type: {type(data_object).__name__}")
    except Exception as e:
        print(f"Failed execution for input type: {type(data_object).__name__}, Error: {e}")

Called add_one_array
Successful execution for input type: ndarray
Called add_one_xarray
Successful execution for input type: DataArray
Called add_one_xarray
Successful execution for input type: Dataset
Failed execution for input type: SimpleFieldList, Error: No dispatcher matched for function add_one with argument t of type <class 'earthkit.data.indexing.simple.SimpleFieldList'>, and no default dispatcher specified.
Failed execution for input type: list, Error: No dispatcher matched for function add_one with argument t of type <class 'list'>, and no default dispatcher specified.

In the offset example, we only accept array objects. This is a strict match, so list, fieldlist and xarray input objects are rejected.

for data_object in [TEST_NUMPY_ARRAY, TEST_XARRAY_DATAARRAY, TEST_XARRAY_DATASET, TEST_FIELDLIST, [1, 2, 3]]:
    try:
        result = offset(data_object)
        print(f"Successful execution for input type: {type(data_object).__name__}")
    except Exception as e:
        print(f"Failed execution for input type: {type(data_object).__name__}, Error: {e}")

Called offset_array
Successful execution for input type: ndarray
Failed execution for input type: DataArray, Error: No dispatcher matched for function offset with argument t of type <class 'xarray.core.dataarray.DataArray'>, and no default dispatcher specified.
Failed execution for input type: Dataset, Error: No dispatcher matched for function offset with argument t of type <class 'xarray.core.dataset.Dataset'>, and no default dispatcher specified.
Failed execution for input type: SimpleFieldList, Error: No dispatcher matched for function offset with argument t of type <class 'earthkit.data.indexing.simple.SimpleFieldList'>, and no default dispatcher specified.
Failed execution for input type: list, Error: No dispatcher matched for function offset with argument t of type <class 'list'>, and no default dispatcher specified.

In the offset_array_like example, we only accept array and array_like objects. This accepts list, fieldlist and xarray.DataArray input objects, and they are all dispatched to the array function. The xarray.Dataset object is still rejected due to an unmatched dispatcher.

for data_object in [TEST_NUMPY_ARRAY, TEST_XARRAY_DATAARRAY, TEST_XARRAY_DATASET, TEST_FIELDLIST, [1, 2, 3]]:
    try:
        result = offset_array_like(data_object)
        print(f"Successful execution for input type: {type(data_object).__name__}")
    except Exception as e:
        print(f"Failed execution for input type: {type(data_object).__name__}, Error: {e}")

Called offset_array
Successful execution for input type: ndarray
Called offset_array
Successful execution for input type: DataArray
Failed execution for input type: Dataset, Error: No dispatcher matched for function offset with argument t of type <class 'xarray.core.dataset.Dataset'>, and no default dispatcher specified.
Called offset_array
Successful execution for input type: SimpleFieldList
Called offset_array
Successful execution for input type: list