def do_draw(self, data):
elements = cu.many(
data,
min_size=self.min_dims,
max_size=self.max_dims,
average_size=min(
max(self.min_dims * 2, self.min_dims + 5),
0.5 * (self.min_dims + self.max_dims),
),
)
result = []
reversed_shape = tuple(self.shape[::-1])
while elements.more():
if len(result) < len(self.shape):
# Shrinks towards original shape
if reversed_shape[len(result)] == 1:
if self.min_side <= 1 and not data.draw(st.booleans()):
side = 1
else:
side = data.draw(self.side_strat)
elif self.max_side >= reversed_shape[len(result)] and (
not self.min_side <= 1 <= self.max_side or data.draw(st.booleans())
):
side = reversed_shape[len(result)]
else:
side = 1
else:
side = data.draw(self.side_strat)
result.append(side)
assert self.min_dims <= len(result) <= self.max_dims
assert all(self.min_side <= s <= self.max_side for s in result)
return tuple(reversed(result))
def from_dtype(dtype):
# type: (np.dtype) -> st.SearchStrategy[Any]
"""Creates a strategy which can generate any value of the given dtype."""
check_type(np.dtype, dtype, "dtype")
# Compound datatypes, eg 'f4,f4,f4'
if dtype.names is not None:
# mapping np.void.type over a strategy is nonsense, so return now.
return st.tuples(*[from_dtype(dtype.fields[name][0]) for name in dtype.names])
# Subarray datatypes, eg '(2, 3)i4'
if dtype.subdtype is not None:
subtype, shape = dtype.subdtype
return arrays(subtype, shape)
# Scalar datatypes
if dtype.kind == u"b":
result = st.booleans() # type: SearchStrategy[Any]
elif dtype.kind == u"f":
if dtype.itemsize == 2:
result = st.floats(width=16)
elif dtype.itemsize == 4:
result = st.floats(width=32)
else:
result = st.floats()
elif dtype.kind == u"c":
if dtype.itemsize == 8:
float32 = st.floats(width=32)
result = st.builds(complex, float32, float32)
else:
result = st.complex_numbers()
elif dtype.kind in (u"S", u"a"):
# Numpy strings are null-terminated; only allow round-trippable values.
# `itemsize == 0` means 'fixed length determined at array creation'
result = st.binary(max_size=dtype.itemsize or None).filter(
lambda b: b[-1:] != b"\0"
)
elif dtype.kind == u"u":
result = st.integers(min_value=0, max_value=2 ** (8 * dtype.itemsize) - 1)
elif dtype.kind == u"i":
overflow = 2 ** (8 * dtype.itemsize - 1)
result = st.integers(min_value=-overflow, max_value=overflow - 1)
elif dtype.kind == u"U":
# Encoded in UTF-32 (four bytes/codepoint) and null-terminated
result = st.text(max_size=(dtype.itemsize or 0) // 4 or None).filter(
lambda b: b[-1:] != u"\0"
)
elif dtype.kind in (u"m", u"M"):
if "[" in dtype.str:
res = st.just(dtype.str.split("[")[-1][:-1])
else:
res = st.sampled_from(TIME_RESOLUTIONS)
result = st.builds(dtype.type, st.integers(-(2 ** 63), 2 ** 63 - 1), res)
else:
raise InvalidArgument(u"No strategy inference for {}".format(dtype))
return result.map(dtype.type)
def all_types(draw):
return draw(
one_of(
text(),
integers(),
none(),
booleans(),
floats(),
tuples(),
times(),
uuids(),
lists(integers()),
dictionaries(text(), text()),
))
def do_draw(self, data):
# General plan: determine the actual selection up front with a straightforward
# approach that shrinks well, then complicate it by inserting other things.
result = []
for dim_size in self.shape:
if dim_size == 0:
result.append(slice(None))
continue
strategy = st.integers(-dim_size, dim_size - 1) | st.slices(dim_size)
result.append(data.draw(strategy))
# Insert some number of new size-one dimensions if allowed
result_dims = sum(isinstance(idx, slice) for idx in result)
while (
self.allow_newaxis
and result_dims < self.max_dims
and (result_dims < self.min_dims or data.draw(st.booleans()))
):
result.insert(data.draw(st.integers(0, len(result))), np.newaxis)
result_dims += 1
# Check that we'll have the right number of dimensions; reject if not.
# It's easy to do this by construction iff you don't care about shrinking,
# which is really important for array shapes. So we filter instead.
assume(self.min_dims <= result_dims <= self.max_dims)
# This is a quick-and-dirty way to insert ..., xor shorten the indexer,
# but it means we don't have to do any structural analysis.
if self.allow_ellipsis and data.draw(st.booleans()):
# Choose an index; then replace all adjacent whole-dimension slices.
i = j = data.draw(st.integers(0, len(result)))
while i > 0 and result[i - 1] == slice(None):
i -= 1
while j < len(result) and result[j] == slice(None):
j += 1
result[i:j] = [Ellipsis]
else:
while result[-1:] == [slice(None, None)] and data.draw(st.integers(0, 7)):
result.pop()
return tuple(result)
def comprehension(draw, expression) -> ast.comprehension:
return ast.comprehension(target=draw(Name(ast.Store)),
iter=draw(expression),
ifs=draw(lists(expression, min_size=0,
max_size=3)),
is_async=draw(booleans()))
