import datetime
from os import environ
import hypothesis as hyp
import hypothesis.strategies as st
from hypothesis.control import current_build_context
from ehrql.query_model.nodes import (
AggregateByPatient,
Case,
Dataset,
Filter,
Function,
InlinePatientTable,
PickOneRowPerPatient,
Position,
SelectColumn,
SelectPatientTable,
SelectTable,
SeriesCollectionFrame,
Sort,
Value,
)
from ehrql.query_model.population_validation import (
ValidationError,
validate_population_definition,
)
from .generic_strategies import usually
from .ignored_errors import get_ignored_error_type
# Max depth
#
# There are various points at which we generate deeply recursive data
# which hits Hypothesis's recursion limits, and we need to stop going deeper
# at this point and force generating a terminating node.
#
# Otherwise, the generated graph can continue forever, and will eventually hit the
# hypothesis limit (100) and will be abandoned. This results in too many invalid examples,
# which triggers the too-many-filters healthcheck.
#
# If the max limit is set high - e.g. if we always let it go to 100 and then return our
# default terminating node, generating the examples takes a really long time. Setting it
# too low means that hypothesis takes too long to shrink examples.
#
# The default is therefore set, somewhat arbitrarily, to 15.
MAX_DEPTH = int(environ.get("GENTEST_MAX_DEPTH", 15))
def depth_exceeded():
ctx = current_build_context()
return ctx.data.depth > MAX_DEPTH
@st.composite
def _should_stop(draw):
"""Returns True if we need to stop and generate a terminating node."""
# Generally speaking we want this to return False unless it needs
# to return True. This need can either come from the fact that
# we've exceeded the maximum depth, or because the shrinker told
# us to.
#
# In the former case, we still need to draw a variable that says
# we should, because this gives us the shrinker the opportunity to
# set that decision to false, which makes us no longer dependent on
# hitting the maximum depth to generate a terminating node here.
should_continue = draw(usually)
if depth_exceeded():
should_continue = False
return not should_continue
should_stop = _should_stop()
@st.composite
def depth_bounded_one_of(draw, *options):
"""Equivalent to `one_of` but if we've got too deep always uses the first option."""
assert options
# Similar to how `should_stop` works, we always draw the choice, but if
# we've exceeded the current maximum depth, we pretend that we got a zero
# even if we didn't. When the shrinker runs it will change this to zero
# for real, and then we no longer need to hit maximum depth for this branch
# to trigger.
i = draw(st.integers(0, len(options) - 1))
if depth_exceeded():
i = 0
return draw(options[i])
# This module defines a set of recursive Hypothesis strategies for generating query model graphs.
#
# There are a few points where we deliberate order the types that we choose from, with the
# "simplest" first (by some subjective measure). This is to enable Hypothesis to more effectively
# explore the query space and to "shrink" examples when it finds errors. These points are commented
# below.
#
# We use several Hypothesis combinators for defining our strategies. Most (`one_of`, `just`,
# `sampled_from`) are fairly self-explanatory. A couple are worth clarifying.
# * `st.builds()` is used to construct objects, it takes the class and strategies
# corresponding to the constructor arguments.
# * `@st.composite` allows us to define a strategy by composing other strategies with
# arbitrary Python code; it adds a `draw` argument which is part of the machinery that
# enables this composition but which doesn't form part of the signature of the resulting
# strategy function.
def dataset(patient_tables, event_tables, schema, value_strategies):
# Every inner-function here returns a Hypothesis strategy for creating the thing it is named
# for, not the thing itself.
#
# Several of these strategy functions ignore one or more of their arguments in order to make
# them uniform with other functions that return the same sort of strategy. Such ignored
# arguments are named with a leading underscore.
# Series strategies
#
# Whenever a series is needed, we call series() passing the type of the series and frame that
# it should be built on (these are either constrained by the context in which the series is to
# be used or chosen arbitrarily by the caller).
#
# This strategy then chooses an arbitrary concrete series that respects the constraints imposed
# by the passed type and frame.
#
# A note on frames and domains:
#
# When we pass `frame` as an argument to a series strategy function, the intended semantics
# are always "construct a series that is _consistent_ with this frame". It's always
# permitted to return a one-row-per-patient series, because such series can always be
# composed a many-rows-per-patient series; so there are series strategy functions that,
# always or sometimes, ignore the frame argument.
COMPARABLE_TYPES = [t for t in value_strategies.keys() if t is not bool]
@st.composite
def series(draw, type_, frame):
if draw(should_stop): # pragma: no cover
return draw(select_column(type_, frame))
class DomainConstraint:
PATIENT = (True,)
NON_PATIENT = (False,)
ANY = (True, False)
# Order matters: "simpler" first (see header comment)
series_constraints = {
select_column: (value_strategies.keys(), DomainConstraint.ANY),
exists: ({bool}, DomainConstraint.PATIENT),
count: ({int}, DomainConstraint.PATIENT),
count_distinct: ({int}, DomainConstraint.PATIENT),
min_: (COMPARABLE_TYPES, DomainConstraint.PATIENT),
max_: (COMPARABLE_TYPES, DomainConstraint.PATIENT),
sum_: ({int, float}, DomainConstraint.PATIENT),
mean: ({float}, DomainConstraint.PATIENT),
is_null: ({bool}, DomainConstraint.ANY),
not_: ({bool}, DomainConstraint.ANY),
year_from_date: ({int}, DomainConstraint.ANY),
month_from_date: ({int}, DomainConstraint.ANY),
day_from_date: ({int}, DomainConstraint.ANY),
to_first_of_year: ({datetime.date}, DomainConstraint.ANY),
to_first_of_month: ({datetime.date}, DomainConstraint.ANY),
cast_to_float: ({float}, DomainConstraint.ANY),
cast_to_int: ({int}, DomainConstraint.ANY),
negate: ({int, float}, DomainConstraint.ANY),
eq: ({bool}, DomainConstraint.ANY),
ne: ({bool}, DomainConstraint.ANY),
string_contains: ({bool}, DomainConstraint.ANY),
in_: ({bool}, DomainConstraint.ANY),
and_: ({bool}, DomainConstraint.ANY),
or_: ({bool}, DomainConstraint.ANY),
lt: ({bool}, DomainConstraint.ANY),
gt: ({bool}, DomainConstraint.ANY),
le: ({bool}, DomainConstraint.ANY),
ge: ({bool}, DomainConstraint.ANY),
add: ({int, float}, DomainConstraint.ANY),
subtract: ({int, float}, DomainConstraint.ANY),
multiply: ({int, float}, DomainConstraint.ANY),
truediv: ({float}, DomainConstraint.ANY),
floordiv: ({int}, DomainConstraint.ANY),
date_add_years: ({datetime.date}, DomainConstraint.ANY),
date_add_months: ({datetime.date}, DomainConstraint.ANY),
date_add_days: ({datetime.date}, DomainConstraint.ANY),
date_difference_in_years: ({int}, DomainConstraint.ANY),
date_difference_in_months: ({int}, DomainConstraint.ANY),
date_difference_in_days: ({int}, DomainConstraint.ANY),
count_episodes: ({int}, DomainConstraint.PATIENT),
case: ({int, float, bool, datetime.date}, DomainConstraint.ANY),
maximum_of: (COMPARABLE_TYPES, DomainConstraint.ANY),
minimum_of: (COMPARABLE_TYPES, DomainConstraint.ANY),
}
series_types = series_constraints.keys()
def constraints_match(series_type):
type_constraint, domain_constraint = series_constraints[series_type]
return (
type_ in type_constraint
and is_one_row_per_patient_frame(frame) in domain_constraint
)
possible_series = [s for s in series_types if constraints_match(s)]
assert possible_series, f"No series matches {type_}, {type(frame)}"
series_strategy = draw(st.sampled_from(possible_series))
return draw(series_strategy(type_, frame))
def value(type_, _frame):
return st.builds(Value, value_strategies[type_])
def select_column(type_, frame):
column_names = [n for n, t in schema.column_types if t == type_]
return st.builds(SelectColumn, st.just(frame), st.sampled_from(column_names))
def exists(_type, _frame):
return st.builds(AggregateByPatient.Exists, any_frame())
def count(_type, _frame):
return st.builds(AggregateByPatient.Count, any_frame())
@st.composite
def count_distinct(draw, _type, _frame):
type_ = draw(any_type())
frame = draw(many_rows_per_patient_frame())
return AggregateByPatient.CountDistinct(draw(series(type_, frame)))
@st.composite
def count_episodes(draw, _type, _frame):
frame = draw(many_rows_per_patient_frame())
date_series = draw(series(datetime.date, frame))
maximum_gap_days = draw(st.integers(1, 5))
return AggregateByPatient.CountEpisodes(date_series, maximum_gap_days)
def min_(type_, _frame):
return aggregation_operation(type_, AggregateByPatient.Min)
def max_(type_, _frame):
return aggregation_operation(type_, AggregateByPatient.Max)
def sum_(type_, _frame):
return aggregation_operation(type_, AggregateByPatient.Sum)
def combine_as_set(type_, _frame):
return aggregation_operation(type_, AggregateByPatient.CombineAsSet)
@st.composite
def mean(draw, _type, _frame):
type_ = draw(any_numeric_type())
frame = draw(many_rows_per_patient_frame())
return AggregateByPatient.Mean(draw(series(type_, frame)))
@st.composite
def aggregation_operation(draw, type_, aggregation):
# An aggregation operation that returns a patient series but takes a
# series drawn from a many-rows-per-patient frame
frame = draw(many_rows_per_patient_frame())
return aggregation(draw(series(type_, frame)))
@st.composite
def is_null(draw, _type, frame):
type_ = draw(any_type())
return Function.IsNull(draw(series(type_, frame)))
def not_(type_, frame):
return st.builds(Function.Not, series(type_, frame))
def year_from_date(_type, frame):
return st.builds(Function.YearFromDate, series(datetime.date, frame))
def month_from_date(_type, frame):
return st.builds(Function.MonthFromDate, series(datetime.date, frame))
def day_from_date(_type, frame):
return st.builds(Function.DayFromDate, series(datetime.date, frame))
def to_first_of_year(_type, frame):
return st.builds(Function.ToFirstOfYear, series(datetime.date, frame))
def to_first_of_month(_type, frame):
return st.builds(Function.ToFirstOfMonth, series(datetime.date, frame))
@st.composite
def cast_to_float(draw, _type, frame):
type_ = draw(any_numeric_type())
return Function.CastToFloat(draw(series(type_, frame)))
@st.composite
def cast_to_int(draw, type_, frame):
type_ = draw(any_numeric_type())
return Function.CastToInt(draw(series(type_, frame)))
def negate(type_, frame):
return st.builds(Function.Negate, series(type_, frame))
@st.composite
def eq(draw, _type, frame):
type_ = draw(any_type())
return draw(binary_operation(type_, frame, Function.EQ))
@st.composite
def ne(draw, _type, frame):
type_ = draw(any_type())
return draw(binary_operation(type_, frame, Function.NE))
def string_contains(_type, frame):
return binary_operation(str, frame, Function.StringContains)
@st.composite
def in_(draw, _type, frame):
type_ = draw(any_type())
if not draw(st.booleans()):
rhs = Value(
frozenset(
draw(st.sets(value_strategies[type_], min_size=0, max_size=5))
)
)
else:
rhs = draw(combine_as_set(type_, frame))
return Function.In(draw(series(type_, frame)), rhs)
def and_(type_, frame):
return binary_operation(type_, frame, Function.And, allow_value=False)
def or_(type_, frame):
return binary_operation(type_, frame, Function.Or, allow_value=False)
@st.composite
def lt(draw, _type, frame):
type_ = draw(any_comparable_type())
return draw(binary_operation(type_, frame, Function.LT))
@st.composite
def gt(draw, _type, frame):
type_ = draw(any_comparable_type())
return draw(binary_operation(type_, frame, Function.GT))
@st.composite
def le(draw, _type, frame):
type_ = draw(any_comparable_type())
return draw(binary_operation(type_, frame, Function.LE))
@st.composite
def ge(draw, _type, frame):
type_ = draw(any_comparable_type())
return draw(binary_operation(type_, frame, Function.GE))
def add(type_, frame):
return binary_operation(type_, frame, Function.Add)
def subtract(type_, frame):
return binary_operation(type_, frame, Function.Subtract)
def multiply(type_, frame):
return binary_operation(type_, frame, Function.Multiply)
def truediv(type_, frame):
return binary_operation(type_, frame, Function.TrueDivide)
def floordiv(type_, frame):
return binary_operation(type_, frame, Function.FloorDivide)
def date_add_years(type_, frame):
return binary_operation_with_types(type_, int, frame, Function.DateAddYears)
def date_add_months(type_, frame):
return binary_operation_with_types(type_, int, frame, Function.DateAddMonths)
def date_add_days(type_, frame):
return binary_operation_with_types(type_, int, frame, Function.DateAddDays)
def date_difference_in_years(type_, frame):
return binary_operation(datetime.date, frame, Function.DateDifferenceInYears)
def date_difference_in_months(type_, frame):
return binary_operation(datetime.date, frame, Function.DateDifferenceInMonths)
def date_difference_in_days(type_, frame):
return binary_operation(datetime.date, frame, Function.DateDifferenceInDays)
@st.composite
def case(draw, type_, frame):
# case takes a mapping argument which is a dict where:
# - keys are a bool series
# - values are either a series or Value of `type_` or None
# It also takes a default, which can be None or a Value or series of `type_`
key_st = series(bool, frame)
value_st = st.one_of(st.none(), value(type_, frame), series(type_, frame))
mapping_st = st.dictionaries(key_st, value_st, min_size=1, max_size=3)
default_st = st.one_of(st.none(), value(type_, frame), series(type_, frame))
mapping = draw(mapping_st)
default = draw(default_st)
# A valid Case needs at least one non-NULL value or a default
hyp.assume(not all(v is None for v in [default, *mapping.values()]))
return Case(mapping, default)
def binary_operation(type_, frame, operator_func, allow_value=True):
# A strategy for operations that take lhs and rhs arguments of the
# same type
return binary_operation_with_types(
type_, type_, frame, operator_func, allow_value=allow_value
)
@st.composite
def binary_operation_with_types(
draw, lhs_type, rhs_type, frame, operator_func, allow_value=True
):
# A strategy for operations that take lhs and rhs arguments with specified lhs
# and rhs types (which may be different)
# A binary operation has 2 inputs, which are
# 1) A series drawn from the specified frame
# 2) one of:
# a) A series drawn from the specified frame
# b) A series drawn from any one-row-per-patient-frame
# c) A series that is a Value
# For certain operations, Value is not allowed; Specifically, for boolean operations
# i.e. and/or which take two boolean series as inputs, we exclude operations that would
# use True/False constant Values. These are unlikely to be seen in the wild, and cause
# particularly nonsensical Case statements in generative test examples.
# first pick an "other" input series (i.e. #2 above), either a value series (if allowed)
# or a series drawn from a frame
series_options = [value, series] if allow_value else [series]
other_series = draw(st.sampled_from(series_options))
# Now pick a frame for the series to be drawn from
# The other frame will either be a new one-row-per-patient-frame or this frame
# (Note if the other_series is a value, the frame will be ignored)
other_frame = draw(st.one_of(one_row_per_patient_frame(), st.just(frame)))
# Pick the order of the lhs and rhs inputs built from the two frames and
# associated strategies
lhs_frame, lhs_input, rhs_frame, rhs_input = draw(
st.sampled_from(
[
(frame, series, other_frame, other_series),
(other_frame, other_series, frame, series),
]
)
)
lhs = draw(lhs_input(lhs_type, lhs_frame))
rhs = draw(rhs_input(rhs_type, rhs_frame))
return operator_func(lhs, rhs)
@st.composite
def nary_operation_with_types(draw, frame, operator_func, series_type):
# A strategy for operations that take _n_ arguments which are expected to be
# the same type
# Decide how many arguments we want – we're intending to test the logic of the
# query engines, not their scaling properties so we don't need too many
num_args = draw(st.integers(1, 4))
# Pick out some arguments (identified by index) to be drawn from other frames
other_frame_args = draw(
st.lists(
# Draw a list of argument indices
st.integers(0, num_args - 1),
# Always leaving at least one argument to be drawn from the original
# frame
max_size=num_args - 1,
unique=True,
)
)
args = []
# Clauses below arranged in order of simplicity (as Hypothesis sees it)
for i in range(num_args):
if i not in other_frame_args:
arg = draw(series(series_type, frame))
else:
# If it's not drawn from the supplied frame then it should be either a
# value or a one-row-per-patient series
if not draw(st.booleans()):
arg = draw(value(series_type, None))
else:
arg = draw(series(series_type, draw(one_row_per_patient_frame())))
args.append(arg)
return operator_func(tuple(args))
def maximum_of(type_, frame):
return nary_operation_with_types(frame, Function.MaximumOf, type_)
def minimum_of(type_, frame):
return nary_operation_with_types(frame, Function.MinimumOf, type_)
def any_type():
return st.sampled_from(list(value_strategies.keys()))
def any_numeric_type():
return st.sampled_from([int, float])
def any_comparable_type():
return st.sampled_from(COMPARABLE_TYPES)
# Frame strategies
#
# The main concern when choosing a frame is whether it has one or many rows per patient. Some
# callers require one or the other, some don't mind; so we provide strategies for each case.
# And sometimes callers need _either_ the frame they have in their hand _or_ an arbitrary
# patient frame, so we provide a strategy for that too.
#
# At variance with the general approach here, many-rows-per-patient frames are created by
# imperatively building stacks of filters on top of select nodes, rather than relying on
# recursion, because it enormously simplifies the logic needed to keep filter conditions
# consistent with the source.
def any_frame():
# Order matters: "simpler" first (see header comment)
return st.one_of(
one_row_per_patient_frame(),
many_rows_per_patient_frame(),
)
def one_row_per_patient_frame():
return depth_bounded_one_of(
select_patient_table(),
pick_one_row_per_patient_frame(),
inline_patient_table(),
)
def many_rows_per_patient_frame():
return depth_bounded_one_of(select_table(), filtered_table())
@st.composite
def filtered_table(draw):
source = draw(select_table())
for _ in range(draw(st.integers(min_value=1, max_value=6))):
source = draw(filter_(source))
return source
@st.composite
def sorted_frame(draw):
# Decide how many Sorts and Filters (if any) we're going to apply
operations = draw(
st.lists(st.sampled_from([sort, filter_]), min_size=1, max_size=9).filter(
lambda ls: (1 <= ls.count(sort) <= 3) and (ls.count(filter_) <= 6)
)
)
# Pick a table and apply the operations
source = draw(select_table())
for operation in operations:
source = draw(operation(source))
return source
@st.composite
def pick_one_row_per_patient_frame(draw):
source = draw(sorted_frame())
sort_order = draw(st.sampled_from([Position.FIRST, Position.LAST]))
return PickOneRowPerPatient(source, sort_order)
def select_table():
return st.builds(SelectTable, st.sampled_from(event_tables), st.just(schema))
def select_patient_table():
return st.builds(
SelectPatientTable, st.sampled_from(patient_tables), st.just(schema)
)
@st.composite
def inline_patient_table(draw):
return InlinePatientTable(
rows=tuple(
draw(
st.lists(
st.tuples(
st.integers(1, 10),
*[
value_strategies[type_]
for name, type_ in schema.column_types
],
),
unique_by=lambda r: r[0],
),
)
),
schema=schema,
)
@st.composite
def filter_(draw, source):
condition = draw(series(bool, draw(ancestor_of(source))))
return Filter(source, condition)
@st.composite
def sort(draw, source):
type_ = draw(any_comparable_type())
sort_by = draw(series(type_, draw(ancestor_of(source))))
return Sort(source, sort_by)
@st.composite
def ancestor_of(draw, frame):
for _ in range(draw(st.integers(min_value=0, max_value=3))):
if hasattr(frame, "source"):
frame = frame.source
else:
break
return frame
# Variable strategy
#
# Puts everything above together to create a variable.
@st.composite
def valid_patient_variable(draw):
type_ = draw(any_type())
frame = draw(one_row_per_patient_frame())
return draw(series(type_, frame))
@st.composite
def valid_event_series(draw):
type_ = draw(any_type())
frame = draw(many_rows_per_patient_frame())
return draw(series(type_, frame))
# A population definition is a boolean-typed variable that meets some additional
# criteria enforced by the query model
@st.composite
def valid_population(draw):
frame = draw(one_row_per_patient_frame())
population = draw(series(bool, frame))
hyp.assume(is_valid_population(population))
return population
return st.builds(
make_dataset,
valid_population(),
valid_patient_variable(),
# Event series is optional
st.one_of(st.none(), valid_event_series()),
)
def make_dataset(population, patient_variable, event_series):
return Dataset(
population=population,
variables={"v": patient_variable},
events=(
{
"event_table": SeriesCollectionFrame({"e": event_series}),
}
if event_series is not None
else {}
),
measures=None,
)
def is_valid_population(series):
try:
validate_population_definition(series)
return True
except ValidationError:
return False
except Exception as e: # pragma: no cover
if get_ignored_error_type(e):
return False
raise
def is_one_row_per_patient_frame(frame):
return isinstance(frame, SelectPatientTable | PickOneRowPerPatient)
