Function outer_array_indexers

xarray/testing/strategies.py:591–634

Generate outer array indexers (vectorized/orthogonal indexing). Parameters ---------- draw : callable The Hypothesis draw function (automatically provided by @st.composite). sizes : dict[Hashable, int] Dictionary mapping dimension names to their sizes. min_dims :

(
    draw,
    /,
    *,
    sizes: dict[Hashable, int],
    min_dims: int = 0,
    max_dims: int | None = None,
    max_size: int = 10,
)

@st.composite
def outer_array_indexers(
    draw,
    /,
    *,
    sizes: dict[Hashable, int],
    min_dims: int = 0,
    max_dims: int | None = None,
    max_size: int = 10,
) -> dict[Hashable, np.ndarray]:
    """Generate outer array indexers (vectorized/orthogonal indexing).

    Parameters
    ----------
    draw : callable
        The Hypothesis draw function (automatically provided by @st.composite).
    sizes : dict[Hashable, int]
        Dictionary mapping dimension names to their sizes.
    min_dims : int, optional
        Minimum number of dimensions to index
    max_dims : int or None, optional
        Maximum number of dimensions to index

    Returns
    -------
    sizes : mapping of hashable to np.ndarray
        Indexers as a dict with keys randomly selected from ``sizes.keys()``.
        Values are 1D numpy arrays of integer indices for each dimension.

    See Also
    --------
    hypothesis.extra.numpy.arrays
    """
    selected_dims = draw(unique_subset_of(sizes, min_size=min_dims, max_size=max_dims))
    idxr = {
        dim: draw(
            npst.arrays(
                dtype=np.int64,
                shape=st.integers(min_value=1, max_value=min(size, max_size)),
                elements=st.integers(min_value=-size, max_value=size - 1),
            )
        )
        for dim, size in selected_dims.items()
    }
    return idxr


@st.composite