def fixed_coords_dataarrays()

in hypothesis_gufunc/extra/xr.py [0:0]

• 6 lines of code
• 4 McCabe index (conditional complexity)

def fixed_coords_dataarrays(dims, coords, dtype=DEFAULT_DTYPE, elements=None):
    """Generate a :class:`xarray:xarray.DataArray` with coordinates that are fixed a-priori.

    Parameters
    ----------
    dims : list(str)
        Dimensions we need to generate coordinates for.
    coords : dict(str, list)
        Dictionary mapping dimension name to its coordinate values.
    dtype : type
        Data type for values in the :class:`xarray:xarray.DataArray`. This can be anything understood by
        :func:`hypothesis:hypothesis.extra.numpy.arrays`.
    elements : SearchStrategy or None
        Strategy to fill the elements of the :class:`xarray:xarray.DataArray`. If `None`, a default is selected based
        on `dtype`.

    Returns
    -------
    da : :class:`xarray:xarray.DataArray`
        :class:`xarray:xarray.DataArray` generated with the specified coordinates and elements from the specified
        strategy.
    """
    shape = [len(coords[dd]) for dd in dims]
    data_st = arrays(dtype, shape, elements=elements)
    coords = {dd: cc for dd, cc in coords.items() if dd in dims}
    S = data_st.map(lambda data: xr.DataArray(data, coords=coords, dims=dims))
    return S