feat: Add apply_tailoring function

Author: gegnew

cellengine/resources/gate.py

@@ -1,5 +1,10 @@
 from __future__ import annotations
 import importlib
+from cellengine.utils.types import (
+    ApplyTailoringInsert,
+    ApplyTailoringUpdate,
+)
+from collections import defaultdict
 from math import pi
 from operator import itemgetter
 from typing import Any, Dict, List, Optional, Union, Tuple, overload
@@ -13,7 +18,6 @@
 from numpy import array, mean, stack
 
 import cellengine as ce
-from cellengine.resources.fcs_file import FcsFile
 from cellengine.resources.population import Population
 from cellengine.utils import parse_fcs_file_args
 from cellengine.utils import converter, generate_id, readonly
@@ -138,10 +142,26 @@ def update_gate_family(experiment_id: str, gid: str, body: Dict) -> None:
         if res["nModified"] < 1:
             raise Warning("No gates updated.")
 
-    def tailor_to(self, fcs_file: FcsFile):
-        self.tailored_per_file = True
-        self.fcs_file_id = fcs_file._id
-        self.update()
+    def apply_tailoring(self, fcs_file_ids: List[str]) -> Dict[str, List[Gate]]:
+        """Tailor this gate to a specific FCS file or files."""
+        payload = ce.APIClient().apply_tailoring(self.experiment_id, self, fcs_file_ids)
+
+        ret = defaultdict(list)
+        for k, v in payload.items():
+            if v:
+                if k == "deleted":
+                    [ret[k].append(i["_id"]) for i in v]  # type: ignore
+                else:
+                    [ret[k].append(self._synthesize_gate(i)) for i in v]  # type: ignore
+        return dict(ret)
+
+    def _synthesize_gate(
+        self,
+        payload: Union[ApplyTailoringInsert, ApplyTailoringUpdate],
+    ):
+        gate = self.to_dict()
+        gate.update(payload)
+        return Gate.from_dict(gate)
 
 
 class RectangleGate(Gate):

cellengine/utils/api_client/APIClient.py

@@ -1,4 +1,5 @@
 from __future__ import annotations
+from cellengine.utils.types import ApplyTailoringRes
 from functools import lru_cache
 from getpass import getpass
 import importlib
@@ -603,12 +604,15 @@ def update_gate_family(self, experiment_id, gid, body: dict = None) -> dict:
             json=body,
         )
 
-    def tailor_to(self, experiment_id, gate_id, fcs_file_id):
-        """Tailor a gate to a file."""
-        gate = self.get_gate(experiment_id, gate_id, as_dict=True)
-        gate["tailoredPerFile"] = True
-        gate["fcsFileId"] = fcs_file_id
-        return self.update_entity(experiment_id, gate_id, "gates", gate)
+    def apply_tailoring(
+        self, experiment_id: str, gate: Gate, fcs_file_ids: List[str]
+    ) -> ApplyTailoringRes:
+        """Tailor a gate to a file or files."""
+        return self._post(
+            f"{self.base_url}/experiments/{experiment_id}/gates/applyTailored",
+            params={"gid": gate.gid},
+            json={"gate": gate.to_dict(), "fcsFileIds": fcs_file_ids},
+        )
 
     def get_plot(
         self,

cellengine/utils/types.py

@@ -0,0 +1,19 @@
+from __future__ import annotations
+import sys
+
+if sys.version_info >= (3, 8):
+    from typing import TypedDict
+else:
+    from typing_extensions import TypedDict
+
+ApplyTailoringInsert = TypedDict("inserted", {"_id": str, "fcsFileId": str})
+ApplyTailoringUpdate = TypedDict("updated", {"_id": str, "fcsFileId": str})
+ApplyTailoringDelete = TypedDict("deleted", {"_id": str, "fcsFileId": str})
+ApplyTailoringRes = TypedDict(
+    "ApplyTailoringRes",
+    {
+        "inserted": ApplyTailoringInsert,
+        "updated": ApplyTailoringUpdate,
+        "deleted": ApplyTailoringDelete,
+    },
+)

tests/unit/resources/test_gates.py

@@ -744,3 +744,49 @@ def test_create_split_gate(client, ENDPOINT_BASE, experiment, scalesets, split_g
     ]
     assert split_gate.names == ["my gate (L)", "my gate (R)"]
     assert split_gate.model["labels"] == [[-199.9, 0.916], [0.9, 0.916]]
+
+
+@responses.activate
+def test_apply_tailoring(client, monkeypatch, ENDPOINT_BASE, rectangle_gate, fcs_files):
+    payload = {
+        "updated": [
+            {"_id": "updatedGateId1", "fcsFileId": "fileId1"},
+            {"_id": "updatedGateId2", "fcsFileId": "fileId2"},
+        ],
+        "inserted": [],
+        "deleted": [{"_id": "deletedGateId", "fcsFileId": "fileId"}],
+    }
+
+    responses.add(
+        responses.POST,
+        f"{ENDPOINT_BASE}/experiments/{EXP_ID}/gates",
+        status=201,
+        json=rectangle_gate,
+    )
+    gate = RectangleGate.create(
+        EXP_ID,
+        x_channel="FSC-A",
+        y_channel="FSC-W",
+        name="my gate",
+        x1=60000,
+        x2=200000,
+        y1=75000,
+        y2=215000,
+        create_population=False
+    )
+
+    responses.add(
+        responses.POST,
+        f"{ENDPOINT_BASE}/experiments/{EXP_ID}/gates/applyTailored?gid={gate.gid}",
+        json=payload,
+    )
+
+    res = gate.apply_tailoring([fcs_files[0]["_id"]])
+
+    assert res["deleted"] == ["deletedGateId"]
+    assert isinstance(res["updated"], list)
+    assert len(res["updated"]) == 2
+    assert res["updated"][0]._id == "updatedGateId1"
+    assert res["updated"][1]._id == "updatedGateId2"
+    assert res["updated"][0].fcs_file_id == "fileId1"
+    assert res["updated"][1].fcs_file_id == "fileId2"

tests/unit/utils/test_api_client.py

@@ -7,7 +7,7 @@
 from cellengine.resources.compensation import Compensation
 from cellengine.resources.experiment import Experiment
 from cellengine.resources.fcs_file import FcsFile
-from cellengine.resources.gate import Gate
+from cellengine.resources.gate import Gate, RectangleGate
 from cellengine.resources.plot import Plot
 from cellengine.resources.population import Population
 from cellengine.resources.scaleset import ScaleSet
@@ -317,3 +317,39 @@ def test_clone_experiment_send_correct_body(client, ENDPOINT_BASE, experiments):
 
     client.clone_experiment(EXP_ID, props={"name": "new exp name"})
     assert responses.calls[0].request.body == b'{"name": "new exp name"}'
+
+
+@responses.activate
+def test_apply_tailoring(client, ENDPOINT_BASE, gates):
+    responses.add(
+        responses.POST,
+        f"{ENDPOINT_BASE}/experiments/{EXP_ID}/gates",
+        json=gates[0],
+    )
+    gate = RectangleGate.create(
+        EXP_ID,
+        x_channel="FSC-A",
+        y_channel="FSC-W",
+        name="my gate",
+        x1=60000,
+        x2=200000,
+        y1=75000,
+        y2=215000,
+    )
+    payload = {
+        "updated": [{"_id": "updatedGateId", "fcsFileId": "fileId"}],
+        "inserted": [],
+        "deleted": [{"_id": "deletedGateId", "fcsFileId": "fileId"}],
+    }
+    responses.add(
+        responses.POST,
+        f"{ENDPOINT_BASE}/experiments/{EXP_ID}/gates/applyTailored?gid={gate.gid}",
+        json=payload,
+    )
+
+    res = client.apply_tailoring(EXP_ID, gate, ["fileID"])
+    assert res == {
+        "updated": [{"_id": "updatedGateId", "fcsFileId": "fileId"}],
+        "inserted": [],
+        "deleted": [{"_id": "deletedGateId", "fcsFileId": "fileId"}],
+    }