Merge pull request #2328 from opensafely-core/additional-pop-constrai…

…nts-measures Allow measures to specify dummy data additional population constraint
opensafely-core · Dec 19, 2024 · 784d011 · 784d011
2 parents 53c57d1 + 8bf2037
commit 784d011
Show file tree

Hide file tree

Showing 3 changed files with 49 additions and 16 deletions.
diff --git a/docs/includes/generated_docs/language__measures.md b/docs/includes/generated_docs/language__measures.md
@@ -80,7 +80,7 @@ this method more than once is an error.
 </div>
 
 <div class="attr-heading" id="Measures.configure_dummy_data">
-  <tt><strong>configure_dummy_data</strong>(<em>population_size=10</em>, <em>legacy=False</em>, <em>timeout=60</em>)</tt>
+  <tt><strong>configure_dummy_data</strong>(<em>population_size=10</em>, <em>legacy=False</em>, <em>timeout=60</em>, <em>additional_population_constraint=None</em>)</tt>
   <a class="headerlink" href="#Measures.configure_dummy_data" title="Permanent link">🔗</a>
 </div>
 <div markdown="block" class="indent">
@@ -98,6 +98,19 @@ 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.
+
 ```py
 measures.configure_dummy_data(population_size=10000)
 ```

diff --git a/ehrql/measures/measures.py b/ehrql/measures/measures.py
@@ -292,6 +292,7 @@ def configure_dummy_data(
         population_size=DummyDataConfig.population_size,
         legacy=DummyDataConfig.legacy,
         timeout=DummyDataConfig.timeout,
+        additional_population_constraint=None,
     ):
         """
         Configure the dummy data to be generated.
@@ -308,13 +309,29 @@ def configure_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.
+
         ```py
         measures.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 configure_disclosure_control(self, *, enabled=True):
         """

diff --git a/ehrql/query_language.py b/ehrql/query_language.py
@@ -50,6 +50,21 @@ class DummyDataConfig:
     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:
     """
@@ -149,21 +164,9 @@ def configure_dummy_data(
         self.dummy_data_config.population_size = population_size
         self.dummy_data_config.legacy = legacy
         self.dummy_data_config.timeout = timeout
-        if additional_population_constraint is not None:
-            validate_patient_series_type(
-                additional_population_constraint,
-                types=[bool],
-                context="additional population constraint",
-            )
-            self.dummy_data_config.additional_population_constraint = (
-                additional_population_constraint._qm_node
-            )
-        else:
-            self.dummy_data_config.additional_population_constraint = None
-        if legacy and additional_population_constraint is not None:
-            raise ValueError(
-                "Cannot provide an additional population constraint in legacy mode."
-            )
+        self.dummy_data_config.set_additional_population_constraint(
+            additional_population_constraint
+        )
 
     def __setattr__(self, name, value):
         if name == "population":