import dataclasses
import datetime
import functools
import operator
import re
from collections import ChainMap
from collections.abc import Callable
from pathlib import Path
from typing import Any, Generic, TypeVar, overload
from ehrql.codes import BaseCode, BaseMultiCodeString
from ehrql.file_formats import read_rows
from ehrql.query_model import nodes as qm
from ehrql.query_model.column_specs import get_column_specs_from_schema
from ehrql.query_model.nodes import get_series_type, has_one_row_per_patient
from ehrql.query_model.population_validation import validate_population_definition
from ehrql.utils import date_utils
from ehrql.utils.string_utils import strip_indent
T = TypeVar("T")
CodeT = TypeVar("CodeT", bound=BaseCode)
MultiCodeStringT = TypeVar("MultiCodeStringT", bound=BaseMultiCodeString)
FloatT = TypeVar("FloatT", bound="FloatFunctions")
DateT = TypeVar("DateT", bound="DateFunctions")
IntT = TypeVar("IntT", bound="IntFunctions")
StrT = TypeVar("StrT", bound="StrFunctions")
VALID_ATTRIBUTE_NAME_RE = re.compile(r"^[A-Za-z]+[A-Za-z0-9_]*$")
# This gets populated by the `__init_subclass__` methods of EventSeries and
# PatientSeries. Its structure is:
#
# (<type>, <is_patient_level>): <SeriesClass>
#
# For example:
#
# (bool, False): BoolEventSeries,
# (bool, True): BoolPatientSeries,
#
REGISTERED_TYPES = {}
class Error(Exception):
"""
Used to translate errors from the query model into something more
ehrQL-appropriate
"""
# Pretend this exception is defined in the top-level `ehrql` module: this allows us
# to avoid leaking the internal `query_language` module into the error messages
# without creating circular import problems.
__module__ = "ehrql"
@dataclasses.dataclass
class DummyDataConfig:
population_size: int = 10
legacy: bool = False
timeout: int = 60
additional_population_constraint: "qm.Series[bool] | None" = None
def set_additional_population_constraint(self, additional_population_constraint):
if additional_population_constraint is not None:
validate_patient_series_type(
additional_population_constraint,
types=[bool],
context="additional population constraint",
)
self.additional_population_constraint = (
additional_population_constraint._qm_node
)
if self.legacy and self.additional_population_constraint is not None:
raise ValueError(
"Cannot provide an additional population constraint in legacy mode."
)
class Dataset:
"""
To create a dataset use the [`create_dataset`](#create_dataset) function.
"""
def __init__(self):
# Set attributes with `object.__setattr__` to avoid using the
# `__setattr__` method on this class, which prohibits use of these
# attribute names
object.__setattr__(self, "_variables", {})
object.__setattr__(self, "dummy_data_config", DummyDataConfig())
object.__setattr__(self, "_events", {})
def define_population(self, population_condition):
"""
Define the condition that patients must meet to be included in the Dataset, in
the form of a [boolean patient series](#BoolPatientSeries).
Example usage:
```python
dataset.define_population(patients.date_of_birth < "1990-01-01")
```
For more detail see the how-to guide on [defining
populations](../how-to/define-population.md).
"""
if hasattr(self, "population"):
raise AttributeError(
"define_population() should be called no more than once"
)
validate_patient_series_type(
population_condition,
types=[bool],
context="population definition",
)
try:
validate_population_definition(population_condition._qm_node)
except qm.ValidationError as exc:
raise Error(str(exc)) from None
object.__setattr__(self, "population", population_condition)
def add_column(self, column_name: str, ehrql_query):
"""
Add a column to the dataset.
_column_name_<br>
The name of the new column, as a string.
_ehrql_query_<br>
An ehrQL query that returns one row per patient.
Example usage:
```python
dataset.add_column("age", patients.age_on("2020-01-01"))
```
Using `.add_column` is equivalent to `=` for adding a single column
but can also be used to add multiple columns, for example by iterating
over a dictionary. For more details see the guide on
"[How to assign multiple columns to a dataset programmatically](../how-to/assign-multiple-columns.md)".
"""
setattr(self, column_name, ehrql_query)
def configure_dummy_data(
self,
*,
population_size=DummyDataConfig.population_size,
legacy=DummyDataConfig.legacy,
timeout=DummyDataConfig.timeout,
additional_population_constraint=None,
):
"""
Configure the dummy data to be generated.
_population_size_<br>
Maximum number of patients to generate.
Note that you may get fewer patients than this if the generator runs out of time
– see `timeout` below.
_legacy_<br>
Use legacy dummy data.
_timeout_<br>
Maximum time in seconds to spend generating dummy data.
_additional_population_constraint_<br>
An additional ehrQL query that can be used to constrain the population that will
be selected for dummy data. This is incompatible with legacy mode.
For example, if you wanted to ensure that two dates appear in a particular order in your
dummy data, you could add ``additional_population_constraint = dataset.first_date <
dataset.second_date``.
You can also combine constraints with ``&`` as normal in ehrQL.
E.g. ``additional_population_constraint = patients.sex.is_in(['male', 'female']) & (
patients.age_on(some_date) < 80)`` would give you dummy data consisting of only men
and women who were under the age of 80 on some particular date.
Example usage:
```python
dataset.configure_dummy_data(population_size=10000)
```
"""
self.dummy_data_config.population_size = population_size
self.dummy_data_config.legacy = legacy
self.dummy_data_config.timeout = timeout
self.dummy_data_config.set_additional_population_constraint(
additional_population_constraint
)
def __setattr__(self, name, value):
if name == "population":
raise AttributeError(
"Cannot set variable 'population'; use define_population() instead"
)
_validate_attribute_name(
name, self._variables | self._events, context="variable"
)
validate_patient_series(value, context=f"variable '{name}'")
self._variables[name] = value
def __getattr__(self, name):
# Make this method accessible while hiding it from autocomplete until we make it
# generally available
if name == "add_event_table":
return self._internal
if name in self._variables:
return self._variables[name]
if name in self._events:
return self._events[name]
if name == "population":
raise AttributeError(
"A population has not been defined; define one with define_population()"
)
else:
raise AttributeError(f"Variable '{name}' has not been defined")
# This method ought to be called `add_event_table` but we're deliberately
# obfuscating its name for now
def _internal(self, name, **event_series):
_validate_attribute_name(name, self._variables | self._events, context="table")
self._events[name] = EventTable(self, **event_series)
def _compile(self):
return qm.Dataset(
population=self.population._qm_node,
variables={k: v._qm_node for k, v in self._variables.items()},
events={k: v._qm_node for k, v in self._events.items()},
measures=None,
)
class EventTable:
def __init__(self, dataset, **series):
# Store reference to the parent dataset to aid debug rendering
object.__setattr__(self, "_dataset", dataset)
object.__setattr__(self, "_series", {})
if not series:
raise ValueError("event tables must be defined with at least one column")
for name, value in series.items():
self.add_column(name, value)
def add_column(self, name, value):
_validate_attribute_name(name, self._series, context="column")
validate_ehrql_series(value, context=f"column {name!r}")
try:
qm_node = qm.SeriesCollectionFrame(
{
name: series._qm_node
for name, series in (self._series | {name: value}).items()
}
)
except qm.PatientDomainError:
raise TypeError(
"event tables must have columns with more than one value per patient; "
"for single values per patient use dataset variables"
)
except qm.DomainMismatchError:
raise Error(
"cannot combine series drawn from different tables; "
"create a new event table for these series"
)
self._series[name] = value
object.__setattr__(self, "_qm_node", qm_node)
def __setattr__(self, name, value):
self.add_column(name, value)
def __getattr__(self, name):
return self._series[name]
def _validate_attribute_name(name, defined_names, context):
if name in defined_names:
raise AttributeError(f"'{name}' is already set and cannot be reassigned")
if name in ("patient_id", "population", "dummy_data_config"):
raise AttributeError(f"'{name}' is not an allowed {context} name")
if not VALID_ATTRIBUTE_NAME_RE.match(name):
raise AttributeError(
f"{context} names must start with a letter, and contain only "
f"alphanumeric characters and underscores (you defined a "
f"{context} '{name}')"
)
def create_dataset():
"""
A dataset defines the patients you want to include in your population and the
variables you want to extract for them.
A dataset definition file must define a dataset called `dataset`:
```python
dataset = create_dataset()
```
Add variables to the dataset as attributes:
```python
dataset.age = patients.age_on("2020-01-01")
```
"""
return Dataset()
# BASIC SERIES TYPES
#
@dataclasses.dataclass(frozen=True)
class BaseSeries:
_qm_node: qm.Node
def __hash__(self):
# The issue here is not mutability but the fact that we overload `__eq__` for
# syntatic sugar, which makes these types spectacularly ill-behaved as dict keys
raise TypeError(f"unhashable type: {self.__class__.__name__!r}")
def __bool__(self):
raise TypeError(
"The keywords 'and', 'or', and 'not' cannot be used with ehrQL, please "
"use the operators '&', '|' and '~' instead.\n"
"(You will also see this error if you try use a chained comparison, "
"such as 'a < b < c'.)"
)
@staticmethod
def _cast(value):
# Series have the opportunity to cast arguments to their methods e.g. to convert
# ISO date strings to date objects. By default, this is a no-op.
return value
# These are the basic operations that apply to any series regardless of type or
# dimension
@overload
def __eq__(self: "PatientSeries", other) -> "BoolPatientSeries": ...
@overload
def __eq__(self: "EventSeries", other) -> "BoolEventSeries": ...
def __eq__(self, other):
"""
Return a boolean series comparing each value in this series with its
corresponding value in `other`.
Note that the result of comparing anything with NULL (including NULL itself) is NULL.
Example usage:
```python
patients.sex == "female"
```
"""
other = self._cast(other)
return _apply(qm.Function.EQ, self, other)
@overload
def __ne__(self: "PatientSeries", other) -> "BoolPatientSeries": ...
@overload
def __ne__(self: "EventSeries", other) -> "BoolEventSeries": ...
def __ne__(self, other):
"""
Return the inverse of `==` above.
Note that the same point regarding NULL applies here.
Example usage:
```python
patients.sex != "unknown"
```
"""
other = self._cast(other)
return _apply(qm.Function.NE, self, other)
@overload
def is_null(self: "PatientSeries") -> "BoolPatientSeries": ...
@overload
def is_null(self: "EventSeries") -> "BoolEventSeries": ...
def is_null(self):
"""
Return a boolean series which is True for each NULL value in this
series and False for each non-NULL value.
Example usage:
```python
patients.date_of_death.is_null()
```
"""
return _apply(qm.Function.IsNull, self)
@overload
def is_not_null(self: "PatientSeries") -> "BoolPatientSeries": ...
@overload
def is_not_null(self: "EventSeries") -> "BoolEventSeries": ...
def is_not_null(self):
"""
Return the inverse of `is_null()` above.
Example usage:
```python
patients.date_of_death.is_not_null()
```
"""
return self.is_null().__invert__()
def when_null_then(self: T, other: T) -> T:
"""
Replace any NULL value in this series with the corresponding value in `other`.
Note that `other` must be of the same type as this series.
Example usage:
```python
(patients.date_of_death < "2020-01-01").when_null_then(False)
```
"""
return case(
when(self.is_not_null()).then(self),
otherwise=self._cast(other),
)
@overload
def is_in(self: "PatientSeries", other) -> "BoolPatientSeries": ...
@overload
def is_in(self: "EventSeries", other) -> "BoolEventSeries": ...
def is_in(self, other):
"""
Return a boolean series which is True for each value in this series which is
contained in `other`.
See how to combine `is_in` with a codelist in
[the how-to guide](../how-to/examples.md/#does-each-patient-have-a-clinical-event-matching-a-code-in-a-codelist).
Example usage:
```python
medications.dmd_code.is_in(["39113311000001107", "39113611000001102"])
```
`other` accepts any of the standard "container" types (tuple, list, set, frozenset,
or dict) or another event series.
"""
if isinstance(other, tuple | list | set | frozenset | dict):
# For iterable arguments, apply any necessary casting and convert to the
# immutable Set type required by the query model. We don't accept arbitrary
# iterables here because too many types in Python are iterable and there's
# the potential for confusion amongst the less experienced of our users.
other = frozenset(map(self._cast, other))
return _apply(qm.Function.In, self, other)
elif isinstance(other, EventSeries):
# We have to use `_convert` and `_wrap` by hand here (rather than using
# `_apply` which does this all for us) because we're constructing a
# `CombineAsSet` query model object which doesn't have a representation in
# the query language.
other_as_set = qm.AggregateByPatient.CombineAsSet(_convert(other))
return _wrap(qm.Function.In, _convert(self), other_as_set)
elif isinstance(other, PatientSeries):
raise TypeError(
"Argument must be an EventSeries (i.e. have many values per patient); "
"you supplied a PatientSeries with only one value per patient"
)
else:
# If the argument is not a supported ehrQL type then we'll get an error here
# (including hopefully helpful errors for common mistakes)
_convert(other)
# Otherwise it _is_ a supported type, but probably not of the right
# cardinality
raise TypeError(
f"Invalid type: {type(other).__qualname__}\n"
f"Note `is_in()` usually expects a list of values rather than a single value"
)
@overload
def is_not_in(self: "PatientSeries", other) -> "BoolPatientSeries": ...
@overload
def is_not_in(self: "EventSeries", other) -> "BoolEventSeries": ...
def is_not_in(self, other):
"""
Return the inverse of `is_in()` above.
"""
return self.is_in(other).__invert__()
def map_values(self, mapping, default=None):
"""
Return a new series with _mapping_ applied to each value. _mapping_ should
be a dictionary mapping one set of values to another.
Example usage:
```python
school_year = patients.age_on("2020-09-01").map_values(
{13: "Year 9", 14: "Year 10", 15: "Year 11"},
default="N/A"
)
```
"""
return case(
*[
when(self == from_value).then(to_value)
for from_value, to_value in mapping.items()
],
otherwise=default,
)
class PatientSeries(BaseSeries):
def __init_subclass__(cls, **kwargs):
super().__init_subclass__(**kwargs)
# Register the series using its `_type` attribute
REGISTERED_TYPES[cls._type, True] = cls
class EventSeries(BaseSeries):
def __init_subclass__(cls, **kwargs):
super().__init_subclass__(**kwargs)
# Register the series using its `_type` attribute
REGISTERED_TYPES[cls._type, False] = cls
def count_distinct_for_patient(self) -> "IntPatientSeries":
"""
Return an [integer patient series](#IntPatientSeries) counting the number of
distinct values for each patient in the series (ignoring any NULL values).
Note that if a patient has no values at all in the series the result will
be zero rather than NULL.
Example usage:
```python
medications.dmd_code.count_distinct_for_patient()
```
"""
return _apply(qm.AggregateByPatient.CountDistinct, self)
# BOOLEAN SERIES
#
class BoolFunctions:
def __and__(self: T, other: T) -> T:
"""
Logical AND
Return a boolean series which is True where both this series and `other` are
True, False where either are False, and NULL otherwise.
Example usage:
```python
is_female_and_alive = patients.is_alive_on("2020-01-01") & patients.sex.is_in(["female"])
```
"""
other = self._cast(other)
return _apply(qm.Function.And, self, other)
def __or__(self: T, other: T) -> T:
"""
Logical OR
Return a boolean series which is True where either this series or `other` is
True, False where both are False, and NULL otherwise.
Example usage:
```python
is_alive = patients.date_of_death.is_null() | patients.date_of_death.is_after("2020-01-01")
```
Note that the above example is equivalent to `patients.is_alive_on("2020-01-01")`.
"""
other = self._cast(other)
return _apply(qm.Function.Or, self, other)
def __invert__(self: T) -> T:
"""
Logical NOT
Return a boolean series which is the inverse of this series i.e. where True
becomes False, False becomes True, and NULL stays as NULL.
Example usage:
```python
is_born_outside_period = ~ patients.date_of_birth.is_on_or_between("2020-03-01", "2020-06-30")
```
"""
return _apply(qm.Function.Not, self)
@overload
def as_int(self: "PatientSeries") -> "IntPatientSeries": ...
@overload
def as_int(self: "EventSeries") -> "IntEventSeries": ...
def as_int(self):
"""
Return each value in this Boolean series as 1 (True) or 0 (False).
"""
return _apply(qm.Function.CastToInt, self)
class BoolPatientSeries(BoolFunctions, PatientSeries):
_type = bool
class BoolEventSeries(BoolFunctions, EventSeries):
_type = bool
# METHODS COMMON TO ALL COMPARABLE TYPES
#
class ComparableFunctions:
@overload
def __lt__(self: "PatientSeries", other) -> "BoolPatientSeries": ...
@overload
def __lt__(self: "EventSeries", other) -> "BoolEventSeries": ...
def __lt__(self, other):
"""
Return a boolean series which is True for each value in this series that is
strictly less than its corresponding value in `other` and False otherwise (or NULL
if either value is NULL).
Example usage:
```python
is_underage = patients.age_on("2020-01-01") < 18
```
"""
other = self._cast(other)
return _apply(qm.Function.LT, self, other)
@overload
def __le__(self: "PatientSeries", other) -> "BoolPatientSeries": ...
@overload
def __le__(self: "EventSeries", other) -> "BoolEventSeries": ...
def __le__(self, other):
"""
Return a boolean series which is True for each value in this series that is less
than or equal to its corresponding value in `other` and False otherwise (or NULL
if either value is NULL).
Example usage:
```python
is_underage = patients.age_on("2020-01-01") <= 17
```
"""
other = self._cast(other)
return _apply(qm.Function.LE, self, other)
@overload
def __ge__(self: "PatientSeries", other) -> "BoolPatientSeries": ...
@overload
def __ge__(self: "EventSeries", other) -> "BoolEventSeries": ...
def __ge__(self, other):
"""
Return a boolean series which is True for each value in this series that is
greater than or equal to its corresponding value in `other` and False otherwise
(or NULL if either value is NULL).
Example usage:
```python
is_adult = patients.age_on("2020-01-01") >= 18
```
"""
other = self._cast(other)
return _apply(qm.Function.GE, self, other)
@overload
def __gt__(self: "PatientSeries", other) -> "BoolPatientSeries": ...
@overload
def __gt__(self: "EventSeries", other) -> "BoolEventSeries": ...
def __gt__(self, other):
"""
Return a boolean series which is True for each value in this series that is
strictly greater than its corresponding value in `other` and False otherwise (or
NULL if either value is NULL).
Example usage:
```python
is_adult = patients.age_on("2020-01-01") > 17
```
"""
other = self._cast(other)
return _apply(qm.Function.GT, self, other)
class ComparableAggregations:
@overload
def minimum_for_patient(self: DateT) -> "DatePatientSeries": ...
@overload
def minimum_for_patient(self: StrT) -> "StrPatientSeries": ...
@overload
def minimum_for_patient(self: IntT) -> "IntPatientSeries": ...
@overload
def minimum_for_patient(self: FloatT) -> "FloatPatientSeries": ...
def minimum_for_patient(self):
"""
Return the minimum value in the series for each patient (or NULL if the patient
has no values).
Example usage:
```python
clinical_events.where(...).numeric_value.minimum_for_patient()
```
"""
return _apply(qm.AggregateByPatient.Min, self)
@overload
def maximum_for_patient(self: DateT) -> "DatePatientSeries": ...
@overload
def maximum_for_patient(self: StrT) -> "StrPatientSeries": ...
@overload
def maximum_for_patient(self: IntT) -> "IntPatientSeries": ...
@overload
def maximum_for_patient(self: FloatT) -> "FloatPatientSeries": ...
def maximum_for_patient(self):
"""
Return the maximum value in the series for each patient (or NULL if the patient
has no values).
Example usage:
```python
clinical_events.where(...).numeric_value.maximum_for_patient()
```
"""
return _apply(qm.AggregateByPatient.Max, self)
# STRING SERIES
#
class StrFunctions(ComparableFunctions):
@overload
def contains(self: "PatientSeries", other) -> "BoolPatientSeries": ...
@overload
def contains(self: "EventSeries", other) -> "BoolEventSeries": ...
def contains(self, other):
"""
Return a boolean series which is True for each string in this series which
contains `other` as a sub-string and False otherwise. For NULL values, the
result is NULL.
Example usage:
```python
is_female = patients.sex.contains("fem")
```
`other` can be another string series, in which case corresponding values
are compared. If either value is NULL the result is NULL.
"""
other = self._cast(other)
return _apply(qm.Function.StringContains, self, other)
class StrAggregations(ComparableAggregations):
"Empty for now"
class StrPatientSeries(StrFunctions, PatientSeries):
_type = str
class StrEventSeries(StrFunctions, StrAggregations, EventSeries):
_type = str
# NUMERIC SERIES
#
class NumericFunctions(ComparableFunctions):
@overload
def __add__(self: IntT, other: IntT | int) -> IntT: ...
@overload
def __add__(self: FloatT, other: FloatT | float) -> FloatT: ...
def __add__(self, other):
"""
Return the sum of each corresponding value in this series and `other` (or NULL
if either is NULL).
"""
other = self._cast(other)
return _apply(qm.Function.Add, self, other)
@overload
def __radd__(self: IntT, other: IntT | int) -> IntT: ...
@overload
def __radd__(self: FloatT, other: FloatT | float) -> FloatT: ...
def __radd__(self, other):
return self + other
@overload
def __sub__(self: IntT, other: IntT | int) -> IntT: ...
@overload
def __sub__(self: FloatT, other: FloatT | float) -> FloatT: ...
def __sub__(self, other):
"""
Return each value in this series with its corresponding value in `other`
subtracted (or NULL if either is NULL).
"""
other = self._cast(other)
return _apply(qm.Function.Subtract, self, other)
@overload
def __rsub__(self: IntT, other: IntT | int) -> IntT: ...
@overload
def __rsub__(self: FloatT, other: FloatT | float) -> FloatT: ...
def __rsub__(self, other):
return other + -self
@overload
def __mul__(self: IntT, other: IntT | int) -> IntT: ...
@overload
def __mul__(self: FloatT, other: FloatT | float) -> FloatT: ...
def __mul__(self, other):
"""
Return the product of each corresponding value in this series and `other` (or
NULL if either is NULL).
"""
other = self._cast(other)
return _apply(qm.Function.Multiply, self, other)
@overload
def __rmul__(self: IntT, other: IntT | int) -> IntT: ...
@overload
def __rmul__(self: FloatT, other: FloatT | float) -> FloatT: ...
def __rmul__(self, other):
return self * other
@overload
def __truediv__(self: "PatientSeries", other) -> "FloatPatientSeries": ...
@overload
def __truediv__(self: "EventSeries", other) -> "FloatEventSeries": ...
def __truediv__(self, other):
"""
Return a series with each value in this series divided by its correponding value
in `other` (or NULL if either is NULL).
Note that the result is always if a float even if the inputs are integers.
"""
other = self._cast(other)
return _apply(qm.Function.TrueDivide, self, other)
@overload
def __rtruediv__(self: "PatientSeries", other) -> "FloatPatientSeries": ...
@overload
def __rtruediv__(self: "EventSeries", other) -> "FloatEventSeries": ...
def __rtruediv__(self, other):
return self / other
@overload
def __floordiv__(self: "PatientSeries", other) -> "IntPatientSeries": ...
@overload
def __floordiv__(self: "EventSeries", other) -> "IntEventSeries": ...
def __floordiv__(self, other):
"""
Return a series with each value in this series divided by its correponding value
in `other` and then rounded **down** to the nearest integer value (or NULL if either
is NULL).
Note that the result is always if an integer even if the inputs are floats.
"""
other = self._cast(other)
return _apply(qm.Function.FloorDivide, self, other)
@overload
def __rfloordiv__(self: "PatientSeries", other) -> "IntPatientSeries": ...
@overload
def __rfloordiv__(self: "EventSeries", other) -> "IntEventSeries": ...
def __rfloordiv__(self, other):
return self // other
def __neg__(self: T) -> T:
"""
Return the negation of each value in this series.
"""
return _apply(qm.Function.Negate, self)
@overload
def as_int(self: "PatientSeries") -> "IntPatientSeries": ...
@overload
def as_int(self: "EventSeries") -> "IntEventSeries": ...
def as_int(self):
"""
Return each value in this series rounded down to the nearest integer.
"""
return _apply(qm.Function.CastToInt, self)
@overload
def as_float(self: "PatientSeries") -> "FloatPatientSeries": ...
@overload
def as_float(self: "EventSeries") -> "FloatEventSeries": ...
def as_float(self):
"""
Return each value in this series as a float (e.g. 10 becomes 10.0).
"""
return _apply(qm.Function.CastToFloat, self)
class NumericAggregations(ComparableAggregations):
@overload
def sum_for_patient(self: FloatT) -> "FloatPatientSeries": ...
@overload
def sum_for_patient(self: IntT) -> "IntPatientSeries": ...
def sum_for_patient(self):
"""
Return the sum of all values in the series for each patient.
"""
return _apply(qm.AggregateByPatient.Sum, self)
def mean_for_patient(self) -> "FloatPatientSeries":
"""
Return the arithmetic mean of any non-NULL values in the series for each
patient.
"""
return _apply(qm.AggregateByPatient.Mean, self)
class IntFunctions(NumericFunctions):
"Currently only needed for type hints to easily tell the difference between int and float series"
class IntPatientSeries(IntFunctions, PatientSeries):
_type = int
class IntEventSeries(IntFunctions, NumericAggregations, EventSeries):
_type = int
class FloatFunctions(NumericFunctions):
@staticmethod
def _cast(value):
"""
Casting int literals to floats. We do not support casting to float for IntSeries.
"""
if isinstance(value, int):
return float(value)
return value
class FloatPatientSeries(FloatFunctions, PatientSeries):
_type = float
class FloatEventSeries(FloatFunctions, NumericAggregations, EventSeries):
_type = float
# DATE SERIES
#
def parse_date_if_str(value):
if isinstance(value, str):
# By default, `fromisoformat()` accepts the alternative YYYYMMDD format. We only
# want to allow the hyphenated version so we pre-validate it.
if not re.match(r"^\d{4}-\d{2}-\d{2}$", value):
raise ValueError(f"Dates must be in YYYY-MM-DD format: {value!r}")
try:
return datetime.date.fromisoformat(value)
except ValueError as e:
raise ValueError(f"{e} in {value!r}") from None
else:
return value
def cast_all_arguments(args):
series_args = [arg for arg in args if isinstance(arg, BaseSeries)]
if series_args:
# Choose the first series arbitrarily, the type checker will enforce that they
# are all the same type in any case
cast = series_args[0]._cast
return tuple(map(cast, args))
else:
# This would only be the case if all the arguments were literals, which would be
# an unusual and pointless bit of ehrQL, but we should handle it without error
return args
# This allows us to get type hints for properties by replacing the
# @property decorator with this decorator. Currently only needed for
# ints. We pass the docstring through so that it can appear in the docs
class int_property(Generic[T]):
def __init__(self, getter: Callable[[Any], T]) -> None:
self.__doc__ = getter.__doc__
self.getter = getter
def __set__(self, instance, value): ...
@overload
def __get__(self, obj: PatientSeries, objtype=None) -> "IntPatientSeries": ...
@overload
def __get__(self, obj: EventSeries, objtype=None) -> "IntEventSeries": ...
def __get__(self, obj, objtype=None):
return self.getter(obj)
class DateFunctions(ComparableFunctions):
@staticmethod
def _cast(value):
return parse_date_if_str(value)
@int_property
def year(self):
"""
Return an integer series giving the year of each date in this series.
"""
return _apply(qm.Function.YearFromDate, self)
@int_property
def month(self):
"""
Return an integer series giving the month (1-12) of each date in this series.
"""
return _apply(qm.Function.MonthFromDate, self)
@int_property
def day(self):
"""
Return an integer series giving the day of the month (1-31) of each date in this
series.
"""
return _apply(qm.Function.DayFromDate, self)
def to_first_of_year(self: T) -> T:
"""
Return a date series with each date in this series replaced by the date of the
first day in its corresponding calendar year.
Example usage:
```python
patients.date_of_death.to_first_of_year()
```
"""
return _apply(qm.Function.ToFirstOfYear, self)
def to_first_of_month(self: T) -> T:
"""
Return a date series with each date in this series replaced by the date of the
first day in its corresponding calendar month.
Example usage:
```python
patients.date_of_death.to_first_of_month()
```
"""
return _apply(qm.Function.ToFirstOfMonth, self)
@overload
def is_before(self: PatientSeries, other) -> BoolPatientSeries: ...
@overload
def is_before(self: EventSeries, other) -> BoolEventSeries: ...
def is_before(self, other):
"""
Return a boolean series which is True for each date in this series that is
strictly earlier than its corresponding date in `other` and False otherwise
(or NULL if either value is NULL).
Example usage:
```python
medications.where(medications.date.is_before("2020-04-01"))
```
"""
return self.__lt__(other)
@overload
def is_on_or_before(self: PatientSeries, other) -> BoolPatientSeries: ...
@overload
def is_on_or_before(self: EventSeries, other) -> BoolEventSeries: ...
def is_on_or_before(self, other):
"""
Return a boolean series which is True for each date in this series that is
earlier than or the same as its corresponding value in `other` and False
otherwise (or NULL if either value is NULL).
Example usage:
```python
medications.where(medications.date.is_on_or_before("2020-03-31"))
```
"""
return self.__le__(other)
@overload
def is_after(self: PatientSeries, other) -> BoolPatientSeries: ...
@overload
def is_after(self: EventSeries, other) -> BoolEventSeries: ...
def is_after(self, other):
"""
Return a boolean series which is True for each date in this series that is
strictly later than its corresponding date in `other` and False otherwise
(or NULL if either value is NULL).
Example usage:
```python
medications.where(medications.date.is_after("2020-03-31"))
```
"""
return self.__gt__(other)
@overload
def is_on_or_after(self: PatientSeries, other) -> BoolPatientSeries: ...
@overload
def is_on_or_after(self: EventSeries, other) -> BoolEventSeries: ...
def is_on_or_after(self, other):
"""
Return a boolean series which is True for each date in this series that is later
than or the same as its corresponding value in `other` and False otherwise (or
NULL if either value is NULL).
Example usage:
```python
medications.where(medications.date.is_on_or_after("2020-04-01"))
```
"""
return self.__ge__(other)
@overload
def is_between_but_not_on(self: PatientSeries, start, end) -> BoolPatientSeries: ...
@overload
def is_between_but_not_on(self: EventSeries, start, end) -> BoolEventSeries: ...
def is_between_but_not_on(self, start, end):
"""
Return a boolean series which is True for each date in this series which is
strictly between (i.e. not equal to) the corresponding dates in `start` and `end`,
and False otherwise.
Example usage:
```python
medications.where(medications.date.is_between_but_not_on("2020-03-31", "2021-04-01"))
```
For each trio of dates being compared, if any date is NULL the result is NULL.
"""
return (self > start) & (self < end)
@overload
def is_on_or_between(self: PatientSeries, start, end) -> BoolPatientSeries: ...
@overload
def is_on_or_between(self: EventSeries, start, end) -> BoolEventSeries: ...
def is_on_or_between(self, start, end):
"""
Return a boolean series which is True for each date in this series which is
between or the same as the corresponding dates in `start` and `end`, and
False otherwise.
Example usage:
```python
medications.where(medications.date.is_on_or_between("2020-04-01", "2021-03-31"))
```
For each trio of dates being compared, if any date is NULL the result is NULL.
"""
return (self >= start) & (self <= end)
@overload
def is_during(self: PatientSeries, interval) -> BoolPatientSeries: ...
@overload
def is_during(self: EventSeries, interval) -> BoolEventSeries: ...
def is_during(self, interval):
"""
The same as `is_on_or_between()` above, but allows supplying a start/end date
pair as single argument.
Example usage:
```python
study_period = ("2020-04-01", "2021-03-31")
medications.where(medications.date.is_during(study_period))
```
Also see the docs on using `is_during` with the
[`INTERVAL` placeholder](../explanation/measures.md/#the-interval-placeholder).
"""
start, end = interval
return self.is_on_or_between(start, end)
def __sub__(self, other):
"""
Return a series giving the difference between each date in this series and
`other` (see [`DateDifference`](#DateDifference)).
Example usage:
```python
age_months = (date("2020-01-01") - patients.date_of_birth).months
```
"""
other = self._cast(other)
if isinstance(other, datetime.date | DateEventSeries | DatePatientSeries):
return DateDifference(self, other)
else:
return NotImplemented
def __rsub__(self, other):
other = self._cast(other)
if isinstance(other, datetime.date | DateEventSeries | DatePatientSeries):
return DateDifference(other, self)
else:
return NotImplemented
class DateAggregations(ComparableAggregations):
def count_episodes_for_patient(self, maximum_gap) -> IntPatientSeries:
"""
Counts the number of "episodes" for each patient where dates which are no more
than `maximum_gap` apart are considered part of the same episode. The
`maximum_gap` duration can be specified in [`days()`](#days) or
[`weeks()`](#weeks).
For example, suppose a patient has the following sequence of events:
Event ID | Date
-- | --
A | 2020-01-01
B | 2020-01-04
C | 2020-01-06
D | 2020-01-10
E | 2020-01-12
And suppose we count the episodes here using a maximum gap of three days:
```python
.count_episodes_for_patient(days(3))
```
We will get an episode count of two: events A, B and C are considered as one
episode and events D and E as another.
Note that events A and C are considered part of the same episode even though
they are more than three days apart because event B is no more than three days
apart from both of them. That is, the clock restarts with each new event in an
episode rather than running from the first event in an episode.
"""
if isinstance(maximum_gap, days):
maximum_gap_days = maximum_gap.value
elif isinstance(maximum_gap, weeks):
maximum_gap_days = maximum_gap.value * 7
else:
raise TypeError("`maximum_gap` must be supplied as `days()` or `weeks()`")
if not isinstance(maximum_gap_days, int):
raise ValueError(
f"`maximum_gap` must be a single, fixed number of "
f"{type(maximum_gap).__name__}"
)
return _wrap(
qm.AggregateByPatient.CountEpisodes,
source=self._qm_node,
maximum_gap_days=maximum_gap_days,
)
class DatePatientSeries(DateFunctions, PatientSeries):
_type = datetime.date
class DateEventSeries(DateFunctions, DateAggregations, EventSeries):
_type = datetime.date
# The default dataclass equality method doesn't work here and while we could define our
# own it wouldn't be very useful for this type
@dataclasses.dataclass(eq=False)
class DateDifference:
"""
Represents the difference between two dates or date series (i.e. it is what you
get when you perform subtractions on [DatePatientSeries](#DatePatientSeries.sub)
or [DateEventSeries](#DateEventSeries.sub)).
"""
lhs: datetime.date | DateEventSeries | DatePatientSeries
rhs: datetime.date | DateEventSeries | DatePatientSeries
@property
def days(self):
"""
The value of the date difference in days (can be positive or negative).
"""
return _apply(qm.Function.DateDifferenceInDays, self.lhs, self.rhs)
@property
def weeks(self):
"""
The value of the date difference in whole weeks (can be positive or negative).
"""
return self.days // 7
@property
def months(self):
"""
The value of the date difference in whole calendar months (can be positive or
negative).
"""
return _apply(qm.Function.DateDifferenceInMonths, self.lhs, self.rhs)
@property
def years(self):
"""
The value of the date difference in whole calendar years (can be positive or
negative).
"""
return _apply(qm.Function.DateDifferenceInYears, self.lhs, self.rhs)
@dataclasses.dataclass
class Duration:
value: int | IntEventSeries | IntPatientSeries
def __init_subclass__(cls, **kwargs):
super().__init_subclass__(**kwargs)
assert cls._date_add_static is not None
assert cls._date_add_qm is not None
# The default dataclass equality/inequality methods don't behave correctly here
def __eq__(self, other) -> bool:
"""
Return True if `other` has the same value and units, and False otherwise.
Hence, the result of `weeks(1) == days(7)` will be False.
"""
if other.__class__ is not self.__class__:
return False
return self.value == other.value
def __ne__(self, other) -> bool:
"""
Return the inverse of `==` above.
"""
# We have to apply different inversion logic depending on whether we have a
# boolean or a BoolSeries
is_equal = self == other
if isinstance(is_equal, bool):
return not is_equal
else:
return is_equal.__invert__()
def __add__(self, other: T) -> T:
"""
If `other` is a date or date series, add this duration to `other`
to produce a new date.
If `other` is another duration with the same units, add the two durations
together to produce a new duration.
"""
other = parse_date_if_str(other)
if isinstance(self.value, int) and isinstance(other, datetime.date):
# If both operands are static values we can perform the date arithmetic
# directly ourselves
return self._date_add_static(other, self.value)
elif isinstance(other, datetime.date | DateEventSeries | DatePatientSeries):
# Otherwise we create the appropriate query model construct
return _apply(self._date_add_qm, other, self.value)
elif isinstance(other, self.__class__):
# Durations of the same type can be added together
return self.__class__(self.value + other.value)
else:
# Nothing else is handled
return NotImplemented
def __sub__(self, other: T) -> T:
"""
Subtract `other` from this duration. `other` must be a
duration in the same units.
"""
return self.__add__(other.__neg__())
def __radd__(self, other: T) -> T:
return self.__add__(other)
def __rsub__(self, other: T) -> T:
return self.__neg__().__add__(other)
def __neg__(self: T) -> T:
"""
Invert this duration, i.e. count the duration backwards in time
if it was originally forwards, and vice versa.
"""
return self.__class__(self.value.__neg__())
def starting_on(self, date) -> list[tuple[datetime.date, datetime.date]]:
"""
Return a list of time intervals covering the duration starting on
`date`. Each interval lasts one unit.
Example usage:
```python
weeks(3).starting_on("2000-01-01")
```
The above would return:
```
[
(date(2000, 1, 1), date(2000, 1, 7)),
(date(2000, 1, 8), date(2000, 1, 14)),
(date(2000, 1, 15), date(2000, 1, 21)),
]
```
Useful for generating the `intervals` arguments to [`Measures`](#Measures).
"""
return self._generate_intervals(date, self.value, 1, "starting_on")
def ending_on(self, date) -> list[tuple[datetime.date, datetime.date]]:
"""
Return a list of time intervals covering the duration ending on
`date`. Each interval lasts one unit.
Example usage:
```python
weeks(3).ending_on("2000-01-21")
```
The above would return:
```
[
(date(2000, 1, 1), date(2000, 1, 7)),
(date(2000, 1, 8), date(2000, 1, 14)),
(date(2000, 1, 15), date(2000, 1, 21)),
]
```
Useful for generating the `intervals` arguments to [`Measures`](#Measures).
"""
return self._generate_intervals(date, self.value, -1, "ending_on")
@classmethod
def _generate_intervals(cls, date, value, sign, method_name):
date = parse_date_if_str(date)
if not isinstance(date, datetime.date):
raise TypeError(
f"{cls.__name__}.{method_name}() can only be used with a literal "
f"date, not a date series"
)
if not isinstance(value, int):
raise TypeError(
f"{cls.__name__}.{method_name}() can only be used with a literal "
f"integer value, not an integer series"
)
if value < 0:
raise ValueError(
f"{cls.__name__}.{method_name}() can only be used with positive numbers"
)
return date_utils.generate_intervals(cls._date_add_static, date, value * sign)
class days(Duration):
"""
Represents a duration of time specified in days.
Example usage:
```python
last_medication_date = medications.sort_by(medications.date).last_for_patient().date
start_date = last_medication_date - days(90)
end_date = last_medication_date + days(90)
```
"""
_date_add_static = staticmethod(date_utils.date_add_days)
_date_add_qm = qm.Function.DateAddDays
class weeks(Duration):
"""
Represents a duration of time specified in weeks.
Example usage:
```python
last_medication_date = medications.sort_by(medications.date).last_for_patient().date
start_date = last_medication_date - weeks(12)
end_date = last_medication_date + weeks(12)
```
"""
_date_add_static = staticmethod(date_utils.date_add_weeks)
@staticmethod
def _date_add_qm(date, num_weeks):
num_days = qm.Function.Multiply(num_weeks, qm.Value(7))
return qm.Function.DateAddDays(date, num_days)
class months(Duration):
"""
Represents a duration of time specified in calendar months.
Example usage:
```python
last_medication_date = medications.sort_by(medications.date).last_for_patient().date
start_date = last_medication_date - months(3)
end_date = last_medication_date + months(3)
```
Consider using [`days()`](#days) or [`weeks()`](#weeks) instead -
see the section on [Ambiguous Dates](#ambiguous-dates) for more.
"""
_date_add_static = staticmethod(date_utils.date_add_months)
_date_add_qm = qm.Function.DateAddMonths
class years(Duration):
"""
Represents a duration of time specified in calendar years.
Example usage:
```python
last_medication_date = medications.sort_by(medications.date).last_for_patient().date
start_date = last_medication_date - years(1)
end_date = last_medication_date + years(1)
```
Consider using [`days()`](#days) or [`weeks()`](#weeks) instead -
see the section on [Ambiguous Dates](#ambiguous-dates) for more.
"""
_date_add_static = staticmethod(date_utils.date_add_years)
_date_add_qm = qm.Function.DateAddYears
# CODE SERIES
#
class CodeFunctions:
def _cast(self, value):
if isinstance(value, str):
return self._type(value)
else:
return value
def to_category(self, categorisation, default=None):
"""
An alias for `map_values` which makes the intention clearer when working with
codelists.
For more detail see [`codelist_from_csv()`](#codelist_from_csv) and the
[how-to guide](../how-to/examples.md/#using-codelists-with-category-columns).
"""
return self.map_values(categorisation, default=default)
class CodePatientSeries(CodeFunctions, PatientSeries):
_type = BaseCode
class CodeEventSeries(CodeFunctions, EventSeries):
_type = BaseCode
class MultiCodeStringFunctions:
def _cast(self, value):
code_type = self._type._code_type()
if isinstance(value, code_type):
# The passed code is of the expected type, so can convert to a string
return value._to_primitive_type()
elif isinstance(value, str) and self._type.regex.fullmatch(value):
# A string that matches the regex for this type
return value
else:
raise TypeError(
f"Expecting a {code_type}, or a string prefix of a {code_type}"
)
def __eq__(self, other):
"""
This operation is not allowed because it is unlikely you would want to match the
values in this field with an exact string e.g.
```python
apcs.all_diagnoses == "||I302, K201, J180 || I302, K200, M920"
```
Instead you should use the `contains` or `contains_any_of` methods.
"""
raise TypeError(
"This column contains multiple clinical codes combined together in a single "
"string. If you want to know if a particular code is contained in the string, "
"please use the `contains()` method"
)
def __ne__(self, other):
"""
See above.
"""
raise TypeError(
"This column contains multiple clinical codes combined together in a single "
"string. If you want to know if a particular code is not contained in the string, "
"please use the `contains()` method."
)
def is_in(self, other):
"""
This operation is not allowed. To check for the presence of any codes in
a codelist, please use the `contains_any_of(codelist)` method instead.
"""
raise TypeError(
"You are attempting to use `.is_in()` on a column that contains multiple "
"clinical codes joined together. This is not allowed. If you want to know "
"if the field contains any of the codes from a codelist, then please use "
"`.contains_any_of(codelist)` instead."
)
def is_not_in(self, other):
"""
This operation is not allowed. To check for the absence of all codes in a codelist,
from a column called `column`, please use `~column.contains_any_of(codelist)`.
NB the `contains_any_of(codelist)` will provide any records that contain any of the
codes, which is then negated with the `~` operator.
"""
raise TypeError(
"You are attempting to use `.is_not_in()` on a column that contains multiple "
"clinical codes joined together. This is not allowed. If you want to know "
"if the column does not contain any of the codes from a codelist, then please use "
"`~column.contains_any_of(codelist)` instead."
)
@overload
def contains(self: PatientSeries, code) -> BoolPatientSeries: ...
@overload
def contains(self: EventSeries, code) -> BoolEventSeries: ...
def contains(self, code):
"""
Check if the multi code field contains a specific code string and
return the result as a boolean series. `code` can
either be a string (and prefix matching works so e.g. "N17" in ICD-10
would match all acute renal failure), or a clinical code.
Example usages:
```python
all_diagnoses.contains("N17")
all_diagnoses.contains(ICD10Code("N170"))
```
"""
code = self._cast(code)
return _apply(qm.Function.StringContains, self, code)
@overload
def contains_any_of(self: PatientSeries, codelist) -> BoolPatientSeries: ...
@overload
def contains_any_of(self: EventSeries, codelist) -> BoolEventSeries: ...
def contains_any_of(self, codelist):
"""
Check if any of the codes in `codelist` occur in the multi code field and
return the result as a boolean series.
As with the `contains(code)` method, the codelist can be a mixture of clinical
codes and string prefixes, as seen in the example below.
Example usage:
```python
all_diagnoses.contains([ICD10Code("N170"), "N17"])
```
"""
conditions = [self.contains(code) for code in codelist]
return functools.reduce(operator.or_, conditions)
class MultiCodeStringPatientSeries(MultiCodeStringFunctions, PatientSeries):
_type = BaseMultiCodeString
class MultiCodeStringEventSeries(MultiCodeStringFunctions, EventSeries):
_type = BaseMultiCodeString
# CONVERT QUERY MODEL SERIES TO EHRQL SERIES
#
def _wrap(qm_cls, *args, **kwargs):
"""
Construct a query model series and wrap it in the ehrQL series class appropriate for
its type and dimension
"""
qm_node = _build(qm_cls, *args, **kwargs)
type_ = get_series_type(qm_node)
is_patient_level = has_one_row_per_patient(qm_node)
try:
cls = REGISTERED_TYPES[type_, is_patient_level]
return cls(qm_node)
except KeyError:
# If we don't have a match for exactly this type then we should have one for a
# superclass. In the case where there are multiple matches, we want the narrowest
# match. E.g. for ICD10MultiCodeString which inherits from BaseMultiCodeString,
# which in turn inherits from str, we want to match BaseMultiCodeString as it
# corresponds to the "closest" series match (in this case MultiCodeStringEventSeries
# rather than the more generic StrEventSeries)
matches = [
{"cls": cls, "depth": type_.__mro__.index(target_type)}
for ((target_type, target_dimension), cls) in REGISTERED_TYPES.items()
if issubclass(type_, target_type) and is_patient_level == target_dimension
]
assert matches, f"No matching query language class for {type_}"
matches.sort(key=lambda k: k["depth"])
cls = matches[0]["cls"]
wrapped = cls(qm_node)
wrapped._type = type_
return wrapped
def _build(qm_cls, *args, **kwargs):
"Construct a query model node, translating any errors as appropriate"
try:
return qm_cls(*args, **kwargs)
except qm.InvalidSortError:
raise Error(
"Cannot sort by a constant value"
# Use `from None` to hide the chained exception
) from None
except qm.DomainMismatchError:
hints = (
" * Reduce one series to have only one value per patient by using\n"
" an aggregation like `maximum_for_patient()`.\n\n"
" * Pick a single row for each patient from the table using\n"
" `first_for_patient()`."
)
if qm_cls is qm.Function.EQ:
hints = (
" * Use `x.is_in(y)` instead of `x == y` to check if values from\n"
" one series match any of the patient's values in the other.\n\n"
f"{hints}"
)
raise Error(
"\n"
"Cannot combine series which are drawn from different tables and both\n"
"have more than one value per patient.\n"
"\n"
"To address this, try one of the following:\n"
"\n"
f"{hints}"
# Use `from None` to hide the chained exception
) from None
except qm.TypeValidationError as exc:
# We deliberately omit information about the query model operation and field
# name here because these often don't match what's used in ehrQL and are liable
# to cause confusion
raise TypeError(
f"Expected type '{_format_typespec(exc.expected)}' "
f"but got '{_format_typespec(exc.received)}'"
# Use `from None` to hide the chained exception
) from None
def _format_typespec(typespec):
# At present we don't do anything beyond formatting as a string and then removing
# the module name prefix from "Series". It might be nice to remove mention of
# "Series" entirely here, but that's a task for another day.
return str(typespec).replace(f"{qm.__name__}.{qm.Series.__qualname__}", "Series")
def _apply(qm_cls, *args):
"""
Applies a query model operation `qm_cls` to its arguments which can be either ehrQL
series or static values, returns an ehrQL series
"""
# Convert all arguments into query model nodes
qm_args = map(_convert, args)
# Construct the query model node and wrap it back up in an ehrQL series
return _wrap(qm_cls, *qm_args)
def _convert(arg):
# Pass null values through unchanged
if arg is None:
return None
# Unpack tuple arguments
elif isinstance(arg, tuple):
return tuple(_convert(a) for a in arg)
# If it's an ehrQL series then get the wrapped query model node
elif isinstance(arg, BaseSeries):
return arg._qm_node
# If it's a static value then we need to be put in a query model Value wrapper
elif isinstance(
arg, bool | int | float | datetime.date | str | BaseCode | frozenset
):
return qm.Value(arg)
else:
raise_helpful_error_if_possible(arg)
raise TypeError(f"Not a valid ehrQL type: {arg!r}")
def Parameter(name, type_):
"""
Return a parameter or placeholder series which can be used to construct a query
"template": a structure which can be turned into a query by substituting in concrete
values for any parameters it contains
"""
return _wrap(qm.Parameter, name, type_)
# FRAME TYPES
#
class BaseFrame:
def __init__(self, qm_node):
self._qm_node = qm_node
def _select_column(self, name):
return _wrap(qm.SelectColumn, source=self._qm_node, name=name)
def exists_for_patient(self) -> BoolPatientSeries:
"""
Return a [boolean patient series](#BoolPatientSeries) which is True for each
patient that has a row in this frame and False otherwise.
Example usage:
```python
pratice_registrations.for_patient_on("2020-01-01").exists_for_patient()
```
"""
return _wrap(qm.AggregateByPatient.Exists, source=self._qm_node)
def count_for_patient(self) -> IntPatientSeries:
"""
Return an [integer patient series](#IntPatientSeries) giving the number of rows each
patient has in this frame.
Note that if a patient has no rows at all in the frame the result will be zero
rather than NULL.
Example usage:
```python
clinical_events.where(clinical_events.date.year == 2020).count_for_patient()
```
"""
return _wrap(qm.AggregateByPatient.Count, source=self._qm_node)
class PatientFrame(BaseFrame):
"""
Frame containing at most one row per patient.
"""
class EventFrame(BaseFrame):
"""
Frame which may contain multiple rows per patient.
"""
def where(self, condition):
"""
Return a new frame containing only the rows in this frame for which `condition`
evaluates True.
Note that this excludes any rows for which `condition` is NULL.
Example usage:
```python
clinical_events.where(clinical_events.date >= "2020-01-01")
```
"""
return self.__class__(
qm.Filter(
source=self._qm_node,
condition=_convert(condition),
)
)
def except_where(self, condition):
"""
Return a new frame containing only the rows in this frame for which `condition`
evaluates False or NULL i.e. the exact inverse of the rows included by
`where()`.
Example usage:
```python
practice_registrations.except_where(practice_registrations.end_date < "2020-01-01")
```
Note that `except_where()` is not the same as `where()` with an inverted condition,
as the latter would exclude rows where `condition` is NULL.
"""
return self.__class__(
qm.Filter(
source=self._qm_node,
condition=qm.Function.Or(
lhs=qm.Function.Not(_convert(condition)),
rhs=qm.Function.IsNull(_convert(condition)),
),
)
)
def sort_by(self, *sort_values):
"""
Return a new frame with the rows sorted for each patient, by
each of the supplied `sort_values`.
Where more than one sort value is supplied then the first (i.e. left-most) value
has highest priority and each subsequent sort value will only be used as a
tie-breaker in case of an exact match among previous values.
Note that NULL is considered smaller than any other value, so you may wish to
filter out NULL values before sorting.
Example usage:
```python
clinical_events.sort_by(clinical_events.date, clinical_events.snomedct_code)
```
"""
# Raise helpful error for easy form of mistake
if string_arg := next((v for v in sort_values if isinstance(v, str)), None):
raise TypeError(
f"to sort by a column use a table attribute like "
f"`{self.__class__.__name__}.{string_arg}` rather than the string "
f'"{string_arg}"'
)
qm_node = self._qm_node
# We expect series to be supplied highest priority first and, as the most
# recently applied Sort operation has the highest priority, we need to apply
# them in reverse order
for series in reversed(sort_values):
qm_node = _build(
qm.Sort,
source=qm_node,
sort_by=_convert(series),
)
cls = make_sorted_event_frame_class(self.__class__)
return cls(qm_node)
class SortedEventFrameMethods:
def first_for_patient(self):
"""
Return a PatientFrame containing, for each patient, the first matching row
according to whatever sort order has been applied.
Note that where there are multiple rows tied for first place then the specific
row returned is picked arbitrarily but consistently i.e. you shouldn't depend on
getting any particular result, but the result you do get shouldn't change unless
the data changes.
Example usage:
```python
medications.sort_by(medications.date).first_for_patient()
```
"""
cls = make_patient_frame_class(self.__class__)
return cls(
qm.PickOneRowPerPatient(
position=qm.Position.FIRST,
source=self._qm_node,
)
)
def last_for_patient(self):
"""
Return a PatientFrame containing, for each patient, the last matching row
according to whatever sort order has been applied.
Note that where there are multiple rows tied for last place then the specific
row returned is picked arbitrarily but consistently i.e. you shouldn't depend on
getting any particular result, but the result you do get shouldn't change unless
the data changes.
Example usage:
```python
medications.sort_by(medications.date).last_for_patient()
```
"""
cls = make_patient_frame_class(self.__class__)
return cls(
qm.PickOneRowPerPatient(
position=qm.Position.LAST,
source=self._qm_node,
)
)
@functools.cache
def make_sorted_event_frame_class(cls):
"""
Given a class return a subclass which has the SortedEventFrameMethods
"""
if issubclass(cls, SortedEventFrameMethods):
return cls
else:
return type(cls.__name__, (SortedEventFrameMethods, cls), {})
@functools.cache
def make_patient_frame_class(cls):
"""
Given an EventFrame subclass return a PatientFrame subclass with the same columns as
the original frame
"""
return type(
cls.__name__,
(PatientFrame,),
get_all_series_and_properties_from_class(cls),
)
def get_all_series_from_class(cls):
# Because `Series` is a descriptor we can't access the column objects via class
# attributes without invoking the descriptor: instead, we have to access them using
# `vars()`. But `vars()` only gives us attributes defined directly on the class, not
# inherited ones. So we reproduce the inheritance behaviour using `ChainMap`.
#
# This is _almost_ exactly what `inspect.getmembers_static` does except that returns
# attributes in lexical order whereas we want to return the original definition
# order.
attrs = ChainMap(*[vars(base) for base in cls.__mro__])
return {key: value for key, value in attrs.items() if isinstance(value, Series)}
def get_all_series_and_properties_from_class(cls):
# Repeating the logic above but also capturing items with the @property decorator.
# This is necessary so we can have properties as well as Series on tables. Keeping
# the other function as there are still other uses where we just want the Series.
attrs = ChainMap(*[vars(base) for base in cls.__mro__])
return {
key: value
for key, value in attrs.items()
if isinstance(value, Series | property)
}
# FRAME CONSTRUCTOR ENTRYPOINTS
#
# A class decorator which replaces the class definition with an appropriately configured
# instance of the class. Obviously this is a _bit_ odd, but I think worth it overall.
# Using classes to define tables is (as far as I can tell) the only way to get nice
# autocomplete and type-checking behaviour for column names. But we don't actually want
# these classes accessible anywhere: users should only be interacting with instances of
# the classes, and having the classes themselves in the module namespaces only makes
# autocomplete more confusing and error prone.
def table(cls: type[T]) -> T:
if PatientFrame in cls.__mro__:
qm_class = qm.SelectPatientTable
elif EventFrame in cls.__mro__:
qm_class = qm.SelectTable
else:
raise Error("Schema class must subclass either `PatientFrame` or `EventFrame`")
qm_node = qm_class(
name=cls.__name__,
schema=get_table_schema_from_class(cls),
)
return cls(qm_node)
def get_table_schema_from_class(cls):
# Get all `Series` objects on the class and determine the schema from them
schema = {
series.name: qm.Column(series.type_, constraints=series.constraints)
for series in get_all_series_from_class(cls).values()
}
return qm.TableSchema(**schema)
# Defines a PatientFrame along with the data it contains. Takes a list (or
# any iterable) of row tuples of the form:
#
# (patient_id, column_1_in_schema, column_2_in_schema, ...)
#
def table_from_rows(rows):
def decorator(cls):
if cls.__bases__ != (PatientFrame,):
raise Error("`@table_from_rows` can only be used with `PatientFrame`")
qm_node = qm.InlinePatientTable(
rows=tuple(rows),
schema=get_table_schema_from_class(cls),
)
return cls(qm_node)
return decorator
# Defines a PatientFrame along with the data it contains. Takes a path to
# a file (feather, csv, csv.gz) with rows of the form:
#
# (patient_id, column_1_in_schema, column_2_in_schema, ...)
#
def table_from_file(path):
path = Path(path)
def decorator(cls):
if cls.__bases__ != (PatientFrame,):
raise Error("`@table_from_file` can only be used with `PatientFrame`")
schema = get_table_schema_from_class(cls)
column_specs = get_column_specs_from_schema(schema)
rows = read_rows(path, column_specs)
qm_node = qm.InlinePatientTable(
rows=rows,
schema=get_table_schema_from_class(cls),
)
return cls(qm_node)
return decorator
# A descriptor which will return the appropriate type of series depending on the type of
# frame it belongs to i.e. a PatientSeries subclass for PatientFrames and an EventSeries
# subclass for EventFrames. This lets schema authors use a consistent syntax when
# defining frames of either type.
class Series(Generic[T]):
def __init__(
self,
type_: type[T],
*,
description="",
constraints=(),
required=True,
implementation_notes_to_add_to_description="",
notes_for_implementors="",
):
self.type_ = type_
self.description = strip_indent(description)
self.constraints = constraints
self.required = required
self.implementation_notes_to_add_to_description = strip_indent(
implementation_notes_to_add_to_description
)
self.notes_for_implementors = strip_indent(notes_for_implementors)
def __set_name__(self, owner, name):
self.name = name
@overload
def __get__(
self: "Series[datetime.date]", instance: PatientFrame, owner
) -> "DatePatientSeries": ...
@overload
def __get__(
self: "Series[datetime.date]", instance: EventFrame, owner
) -> DateEventSeries: ...
@overload
def __get__(
self: "Series[CodeT]", instance: PatientFrame, owner
) -> CodePatientSeries: ...
@overload
def __get__(
self: "Series[CodeT]", instance: EventFrame, owner
) -> CodeEventSeries: ...
@overload
def __get__(
self: "Series[MultiCodeStringT]", instance: PatientFrame, owner
) -> MultiCodeStringPatientSeries: ...
@overload
def __get__(
self: "Series[MultiCodeStringT]", instance: EventFrame, owner
) -> MultiCodeStringEventSeries: ...
@overload
def __get__(
self: "Series[bool]", instance: PatientFrame, owner
) -> BoolPatientSeries: ...
@overload
def __get__(
self: "Series[bool]", instance: EventFrame, owner
) -> BoolEventSeries: ...
@overload
def __get__(
self: "Series[str]", instance: PatientFrame, owner
) -> StrPatientSeries: ...
@overload
def __get__(
self: "Series[str]", instance: EventFrame, owner
) -> "StrEventSeries": ...
@overload
def __get__(
self: "Series[int]", instance: PatientFrame, owner
) -> IntPatientSeries: ...
@overload
def __get__(self: "Series[int]", instance: EventFrame, owner) -> IntEventSeries: ...
@overload
def __get__(
self: "Series[float]", instance: PatientFrame, owner
) -> FloatPatientSeries: ...
@overload
def __get__(
self: "Series[float]", instance: EventFrame, owner
) -> FloatEventSeries: ...
def __get__(self, instance, owner):
if instance is None: # pragma: no cover
return self
return instance._select_column(self.name)
def get_tables_from_namespace(namespace):
"""
Yield all ehrQL tables contained in `namespace`
"""
for attr, value in vars(namespace).items():
if isinstance(value, BaseFrame):
yield attr, value
# CASE EXPRESSION FUNCTIONS
#
class when:
def __init__(self, condition):
condition_qm = _convert(condition)
type_ = get_series_type(condition_qm)
if type_ is not bool:
raise TypeError(
f"invalid case condition:\n"
f"Expecting a boolean series, got series of type"
f" '{type_.__qualname__}'",
)
self._condition = condition_qm
def then(self, value):
return WhenThen(self._condition, _convert(value))
class WhenThen:
def __init__(self, condition, value):
self._condition = condition
self._value = value
def otherwise(self, value):
return case(self, otherwise=value)
def case(*when_thens, otherwise=None):
"""
Take a sequence of condition-values of the form:
```python
when(condition).then(value)
```
And evaluate them in order, returning the value of the first condition which
evaluates True. If no condition matches, return the `otherwise` value (or NULL
if no `otherwise` value is specified).
Example usage:
```python
category = case(
when(size < 10).then("small"),
when(size < 20).then("medium"),
when(size >= 20).then("large"),
otherwise="unknown",
)
```
Note that because the conditions are evaluated in order we don't need the condition
for "medium" to specify `(size >= 10) & (size < 20)` because by the time the
condition for "medium" is being evaluated we already know the condition for "small"
is False.
A simpler form is available when there is a single condition. This example:
```python
category = case(
when(size < 15).then("small"),
otherwise="large",
)
```
can be rewritten as:
```python
category = when(size < 15).then("small").otherwise("large")
```
"""
cases = {}
for case in when_thens:
if isinstance(case, when):
raise TypeError(
"`when(...)` clause missing a `.then(...)` value in `case()` expression"
)
elif (
isinstance(case, BaseSeries)
and isinstance(case._qm_node, qm.Case)
and len(case._qm_node.cases) == 1
):
raise TypeError(
"invalid syntax for `otherwise` in `case()` expression, instead of:\n"
"\n"
" case(\n"
" when(...).then(...).otherwise(...)\n"
" )\n"
"\n"
"You should write:\n"
"\n"
" case(\n"
" when(...).then(...),\n"
" otherwise=...\n"
" )\n"
"\n"
)
elif not isinstance(case, WhenThen):
raise TypeError(
"cases must be specified in the form:\n"
"\n"
" when(<CONDITION>).then(<VALUE>)\n"
"\n"
)
elif case._condition in cases:
raise TypeError("duplicated condition in `case()` expression")
else:
cases[case._condition] = case._value
if not cases:
raise TypeError("`case()` expression requires at least one case")
if otherwise is None and all(value is None for value in cases.values()):
raise TypeError("`case()` expression cannot have all `None` values")
return _wrap(qm.Case, cases, default=_convert(otherwise))
# HORIZONTAL AGGREGATION FUNCTIONS
#
# These cast all arguments to the first Series. So if we have a Series as
# the first arg then we know the return type. However, if the first arg is
# not a Series, then we don't know the return type. E.g. the following examples
# are tricky:
# maximum_of(10, 10, clinical_events.numeric_value) - will return FloatEventSeries
# maximum_of("2024-01-01", "2023-01-01", clinical_events.date) - will return DateEventSeries
@overload
def maximum_of(value: IntT, other_value, *other_values) -> IntT: ...
@overload
def maximum_of(value: FloatT, other_value, *other_values) -> FloatT: ...
@overload
def maximum_of(value: DateT, other_value, *other_values) -> DateT: ...
def maximum_of(value, other_value, *other_values) -> int:
"""
Return the maximum value of a collection of Series or Values, disregarding NULLs.
Example usage:
```python
latest_event_date = maximum_of(event_series_1.date, event_series_2.date, "2001-01-01")
```
"""
args = cast_all_arguments((value, other_value, *other_values))
return _apply(qm.Function.MaximumOf, args)
@overload
def minimum_of(value: IntT, other_value, *other_values) -> IntT: ...
@overload
def minimum_of(value: FloatT, other_value, *other_values) -> FloatT: ...
@overload
def minimum_of(value: DateT, other_value, *other_values) -> DateT: ...
def minimum_of(value, other_value, *other_values):
"""
Return the minimum value of a collection of Series or Values, disregarding NULLs.
Example usage:
```python
ealiest_event_date = minimum_of(event_series_1.date, event_series_2.date, "2001-01-01")
```
"""
args = cast_all_arguments((value, other_value, *other_values))
return _apply(qm.Function.MinimumOf, args)
# ERROR HANDLING
#
def raise_helpful_error_if_possible(arg):
if isinstance(arg, BaseFrame):
raise TypeError(
f"Expecting a series but got a frame (`{arg.__class__.__name__}`): "
f"are you missing a column name?"
)
if callable(arg):
raise TypeError(
f"Function referenced but not called: are you missing parentheses on "
f"`{arg.__name__}()`?"
)
if isinstance(arg, when):
raise TypeError(
"Missing `.then(...).otherwise(...)` conditions on a `when(...)` expression"
)
if isinstance(arg, WhenThen):
raise TypeError(
"Missing `.otherwise(...)` condition on a `when(...).then(...)` expression\n"
"Note: you can use `.otherwise(None)` to get NULL values"
)
def validate_ehrql_series(arg, context):
try:
raise_helpful_error_if_possible(arg)
except TypeError as e:
raise TypeError(f"invalid {context}:\n{e})") from None
if not isinstance(arg, BaseSeries):
raise TypeError(
f"invalid {context}:\n"
f"Expecting an ehrQL series, got type '{type(arg).__qualname__}'"
)
def validate_patient_series(arg, context):
validate_ehrql_series(arg, context)
if not isinstance(arg, PatientSeries):
raise TypeError(
f"invalid {context}:\nExpecting a series with only one value per patient"
)
def validate_patient_series_type(arg, types, context):
validate_patient_series(arg, context)
if arg._type not in types:
types_desc = humanize_list_of_types(types)
article = "an" if types_desc[0] in "aeiou" else "a"
raise TypeError(
f"invalid {context}:\n"
f"Expecting {article} {types_desc} series, got series of type"
f" '{arg._type.__qualname__}'",
)
HUMAN_TYPES = {
bool: "boolean",
int: "integer",
}
def humanize_list_of_types(types):
type_names = [HUMAN_TYPES.get(type_, type_.__qualname__) for type_ in types]
initial = ", ".join(type_names[:-1])
return f"{initial} or {type_names[-1]}" if initial else type_names[-1]
def modify_exception(exc):
# This is our chance to modify exceptions which we didn't raise ourselves to make
# them more helpful or add additional context
if operator := _get_operator_error(exc):
exc.add_note(
_format_operator_error_note(operator),
)
return exc
def _get_operator_error(exc):
# Because `and`, `or` and `not` are control-flow primitives in Python they are not
# overridable and so we're forced to use the bitwise operators for logical
# operations. However these have different precedence rules from those governing the
# standard operators and so it's easy to accidentally do the wrong thing. Here we
# identify errors associated with the logical operators so we can add a note trying
# to explain what might have happened.
if not isinstance(exc, TypeError):
return
# Sadly we have to do this via string matching on the exception text
if match := re.match(
r"(unsupported operand type\(s\) for|bad operand type for unary) ([|&~]):",
str(exc),
):
return match.group(2)
def _format_operator_error_note(operator):
if operator == "|":
example_bad = "a == b | x == y"
example_good = "(a == b) | (x == y)"
elif operator == "&":
example_bad = "a == b & x == y"
example_good = "(a == b) & (x == y)"
elif operator == "~":
example_bad = "~ a == b"
example_good = "~ (a == b)"
else:
assert False
return (
f"\n"
f"WARNING: The `{operator}` operator has surprising precedence rules, meaning\n"
"you may need to add more parentheses to get the correct behaviour.\n"
f"\n"
f"For example, instead of writing:\n"
f"\n"
f" {example_bad}\n"
f"\n"
f"You should write:\n"
f"\n"
f" {example_good}"
)