add docker_image selection parameter

servicex/topcp/topcp.py

@@ -52,6 +52,8 @@ class TopCPQuery(QueryStringGenerator):
     """Toggles off the computation of systematics"""
     no_filter: Optional[bool] = False
     """Save all events regardless of analysis filters (still saves the decision)"""
+    image: Optional[str] = None
+    """Docker image to use"""
 
     @pydantic.model_validator(mode="after")
     def no_input_yaml(self):
@@ -85,6 +87,10 @@ def generate_selection_string(self):
             "no_systematics": self.no_systematics,
             "no_filter": self.no_filter,
         }
+
+        if self.image is not None:
+            query["docker_image"] = self.image
+
         return json.dumps(query)
 
     @classmethod

tests/test_topcp_dataset.py

@@ -113,3 +113,21 @@ def test_yaml_serialization():
 def test_no_yaml():
     with pytest.raises(ValueError):
         TopCPQuery()
+
+
+def test_docker_image(tmp_path):
+    reco_file = tmp_path / "reco.yaml"
+    reco_file.write_text(
+        """
+CommonServices:
+  runSystematics: False
+    """
+    )
+
+    docker_image = "my-custom-image:latest"
+    topcp_query = TopCPQuery(reco=reco_file, image=docker_image)
+    query_string = topcp_query.generate_selection_string()
+    query = json.loads(query_string)
+
+    assert "docker_image" in query, "Missing docker_image key"
+    assert query["docker_image"] == docker_image