def array_dtypes(
subtype_strategy: st.SearchStrategy[np.dtype] = scalar_dtypes(),
*,
min_size: int = 1,
max_size: int = 5,
allow_subarrays: bool = False,
) -> st.SearchStrategy[np.dtype]:
"""Return a strategy for generating array (compound) dtypes, with members
drawn from the given subtype strategy."""
order_check("size", 0, min_size, max_size)
# The empty string is replaced by f{idx}; see #1963 for details. Much easier to
# insist that field names be unique and just boost f{idx} strings manually.
field_names = st.integers(0, 127).map("f{}".format) | st.text(min_size=1)
name_titles = st.one_of(
field_names,
st.tuples(field_names, field_names).filter(lambda ns: ns[0] != ns[1]),
)
elements = st.tuples(name_titles, subtype_strategy)
if allow_subarrays:
elements |= st.tuples(name_titles, subtype_strategy,
array_shapes(max_dims=2, max_side=2))
return st.lists(
elements=elements,
min_size=min_size,
max_size=max_size,
unique_by=(
# Deduplicate by both name and title for efficiency before filtering.
# (Field names must be unique, as must titles, and no intersections)
lambda d: d[0] if isinstance(d[0], str) else d[0][0],
lambda d: d[0] if isinstance(d[0], str) else d[0][1],
),
).filter(_no_title_is_name_of_a_titled_field)
def integer_array_indices(
shape: Shape,
*,
result_shape: st.SearchStrategy[Shape] = array_shapes(),
dtype: np.dtype = "int",
) -> st.SearchStrategy[Tuple[np.ndarray, ...]]:
"""Return a search strategy for tuples of integer-arrays that, when used
to index into an array of shape ``shape``, given an array whose shape
was drawn from ``result_shape``.
Examples from this strategy shrink towards the tuple of index-arrays::
len(shape) * (np.zeros(drawn_result_shape, dtype), )
* ``shape`` a tuple of integers that indicates the shape of the array,
whose indices are being generated.
* ``result_shape`` a strategy for generating tuples of integers, which
describe the shape of the resulting index arrays. The default is
:func:`~hypothesis.extra.numpy.array_shapes`. The shape drawn from
this strategy determines the shape of the array that will be produced
when the corresponding example from ``integer_array_indices`` is used
as an index.
* ``dtype`` the integer data type of the generated index-arrays. Negative
integer indices can be generated if a signed integer type is specified.
Recall that an array can be indexed using a tuple of integer-arrays to
access its members in an arbitrary order, producing an array with an
arbitrary shape. For example:
.. code-block:: pycon
>>> from numpy import array
>>> x = array([-0, -1, -2, -3, -4])
>>> ind = (array([[4, 0], [0, 1]]),) # a tuple containing a 2D integer-array
>>> x[ind] # the resulting array is commensurate with the indexing array(s)
array([[-4, 0],
[0, -1]])
Note that this strategy does not accommodate all variations of so-called
'advanced indexing', as prescribed by NumPy's nomenclature. Combinations
of basic and advanced indexes are too complex to usefully define in a
standard strategy; we leave application-specific strategies to the user.
Advanced-boolean indexing can be defined as ``arrays(shape=..., dtype=bool)``,
and is similarly left to the user.
"""
check_type(tuple, shape, "shape")
check_argument(
shape and all(isinstance(x, int) and x > 0 for x in shape),
f"shape={shape!r} must be a non-empty tuple of integers > 0",
)
check_strategy(result_shape, "result_shape")
check_argument(np.issubdtype(dtype, np.integer),
f"dtype={dtype!r} must be an integer dtype")
signed = np.issubdtype(dtype, np.signedinteger)
def array_for(index_shape, size):
return arrays(
dtype=dtype,
shape=index_shape,
elements=st.integers(-size if signed else 0, size - 1),
)
return result_shape.flatmap(lambda index_shape: st.tuples(*(
array_for(index_shape, size) for size in shape)))