def __call__(self, values):
if not values:
raise IndexError('Cannot choose from empty sequence')
result = choice(self.data, values)
with self.build_context.local():
self.choice_count += 1
note('Choice #%d: %r' % (self.choice_count, result))
return result
def do_draw(self, data):
while True:
try:
result = dt.datetime(
year=cu.centered_integer_range(data, self.min_year, self.max_year, 2000),
month=cu.integer_range(data, 1, 12),
day=cu.integer_range(data, 1, 31),
hour=cu.integer_range(data, 0, 24),
minute=cu.integer_range(data, 0, 59),
second=cu.integer_range(data, 0, 59),
microsecond=cu.integer_range(data, 0, 999999),
)
if not self.allow_naive or (self.timezones and cu.boolean(data)):
result = cu.choice(data, self.timezones).localize(result)
return result
except (OverflowError, ValueError):
pass
def do_draw(self, data):
return d.choice(data, self.elements)
def do_filtered_draw(self, data, filter_strategy):
# Set of indices that have been tried so far, so that we never test
# the same element twice during a draw.
known_bad_indices = set()
def check_index(i):
"""Return ``True`` if the element at ``i`` satisfies the filter
condition.
"""
if i in known_bad_indices:
return False
ok = filter_strategy.condition(self.elements[i])
if not ok:
if not known_bad_indices:
filter_strategy.note_retried(data)
known_bad_indices.add(i)
return ok
# Start with ordinary rejection sampling. It's fast if it works, and
# if it doesn't work then it was only a small amount of overhead.
for _ in range(3):
i = cu.integer_range(data, 0, len(self.elements) - 1)
if check_index(i):
return self.elements[i]
# If we've tried all the possible elements, give up now.
max_good_indices = len(self.elements) - len(known_bad_indices)
if not max_good_indices:
return filter_not_satisfied
# Figure out the bit-length of the index that we will write back after
# choosing an allowed element.
write_length = len(self.elements).bit_length()
# Impose an arbitrary cutoff to prevent us from wasting too much time
# on very large element lists.
cutoff = 10000
max_good_indices = min(max_good_indices, cutoff)
# Before building the list of allowed indices, speculatively choose
# one of them. We don't yet know how many allowed indices there will be,
# so this choice might be out-of-bounds, but that's OK.
speculative_index = cu.integer_range(data, 0, max_good_indices - 1)
# Calculate the indices of allowed values, so that we can choose one
# of them at random. But if we encounter the speculatively-chosen one,
# just use that and return immediately.
allowed_indices = []
for i in range(min(len(self.elements), cutoff)):
if check_index(i):
allowed_indices.append(i)
if len(allowed_indices) > speculative_index:
# Early-exit case: We reached the speculative index, so
# we just return the corresponding element.
data.draw_bits(write_length, forced=i)
return self.elements[i]
# The speculative index didn't work out, but at this point we've built
# the complete list of allowed indices, so we can just choose one of
# them.
if allowed_indices:
i = cu.choice(data, allowed_indices)
data.draw_bits(write_length, forced=i)
return self.elements[i]
# If there are no allowed indices, the filter couldn't be satisfied.
return filter_not_satisfied
def do_draw(self, data):
return cu.choice(data, self.elements)
def do_filtered_draw(self, data, filter_strategy):
# Set of indices that have been tried so far, so that we never test
# the same element twice during a draw.
known_bad_indices = set()
def check_index(i):
"""Return ``True`` if the element at ``i`` satisfies the filter
condition.
"""
if i in known_bad_indices:
return False
ok = filter_strategy.condition(self.elements[i])
if not ok:
if not known_bad_indices:
filter_strategy.note_retried(data)
known_bad_indices.add(i)
return ok
# Start with ordinary rejection sampling. It's fast if it works, and
# if it doesn't work then it was only a small amount of overhead.
for _ in hrange(3):
i = d.integer_range(data, 0, len(self.elements) - 1)
if check_index(i):
return self.elements[i]
# If we've tried all the possible elements, give up now.
max_good_indices = len(self.elements) - len(known_bad_indices)
if not max_good_indices:
return filter_not_satisfied
# Figure out the bit-length of the index that we will write back after
# choosing an allowed element.
write_length = bit_length(len(self.elements))
# Impose an arbitrary cutoff to prevent us from wasting too much time
# on very large element lists.
cutoff = 10000
max_good_indices = min(max_good_indices, cutoff)
# Before building the list of allowed indices, speculatively choose
# one of them. We don't yet know how many allowed indices there will be,
# so this choice might be out-of-bounds, but that's OK.
speculative_index = d.integer_range(data, 0, max_good_indices - 1)
# Calculate the indices of allowed values, so that we can choose one
# of them at random. But if we encounter the speculatively-chosen one,
# just use that and return immediately.
allowed_indices = []
for i in hrange(min(len(self.elements), cutoff)):
if check_index(i):
allowed_indices.append(i)
if len(allowed_indices) > speculative_index:
# Early-exit case: We reached the speculative index, so
# we just return the corresponding element.
data.draw_bits(write_length, forced=i)
return self.elements[i]
# The speculative index didn't work out, but at this point we've built
# the complete list of allowed indices, so we can just choose one of
# them.
if allowed_indices:
i = d.choice(data, allowed_indices)
data.draw_bits(write_length, forced=i)
return self.elements[i]
# If there are no allowed indices, the filter couldn't be satisfied.
return filter_not_satisfied
def test_choice():
assert cu.choice(ConjectureData.for_buffer([1]), [1, 2, 3]) == 2
def do_filtered_draw(self, data, filter_strategy):
# Set of indices that have been tried so far, so that we never test
# the same element twice during a draw.
known_bad_indices = set()
# If we're being called via FilteredStrategy, the filter_strategy argument
# might have additional conditions we have to fulfill.
if isinstance(filter_strategy, FilteredStrategy):
conditions = filter_strategy.flat_conditions
else:
conditions = ()
# Start with ordinary rejection sampling. It's fast if it works, and
# if it doesn't work then it was only a small amount of overhead.
for _ in range(3):
i = cu.integer_range(data, 0, len(self.elements) - 1)
if i not in known_bad_indices:
element = self.get_element(i, conditions=conditions)
if element is not filter_not_satisfied:
return element
if not known_bad_indices:
FilteredStrategy.note_retried(self, data)
known_bad_indices.add(i)
# If we've tried all the possible elements, give up now.
max_good_indices = len(self.elements) - len(known_bad_indices)
if not max_good_indices:
return filter_not_satisfied
# Figure out the bit-length of the index that we will write back after
# choosing an allowed element.
write_length = len(self.elements).bit_length()
# Impose an arbitrary cutoff to prevent us from wasting too much time
# on very large element lists.
cutoff = 10000
max_good_indices = min(max_good_indices, cutoff)
# Before building the list of allowed indices, speculatively choose
# one of them. We don't yet know how many allowed indices there will be,
# so this choice might be out-of-bounds, but that's OK.
speculative_index = cu.integer_range(data, 0, max_good_indices - 1)
# Calculate the indices of allowed values, so that we can choose one
# of them at random. But if we encounter the speculatively-chosen one,
# just use that and return immediately. Note that we also track the
# allowed elements, in case of .map(some_stateful_function)
allowed = []
for i in range(min(len(self.elements), cutoff)):
if i not in known_bad_indices:
element = self.get_element(i, conditions=conditions)
if element is not filter_not_satisfied:
allowed.append((i, element))
if len(allowed) > speculative_index:
# Early-exit case: We reached the speculative index, so
# we just return the corresponding element.
data.draw_bits(write_length, forced=i)
return element
# The speculative index didn't work out, but at this point we've built
# and can choose from the complete list of allowed indices and elements.
if allowed:
i, element = cu.choice(data, allowed)
data.draw_bits(write_length, forced=i)
return element
# If there are no allowed indices, the filter couldn't be satisfied.
return filter_not_satisfied