allow removing custom sample gates (closes #221)

src/flowkit/_models/gate_node.py

@@ -36,8 +36,9 @@ def add_custom_gate(self, sample_id, gate):
         type, have the same gate name, and the same dimension IDs as the
         GateNode's template gate.
 
-        :param sample_id: text string used to identify the custom gate,
-        :param gate: a Gate instance
+        :param sample_id: text string used to identify the custom gate
+        :param gate: a Gate instance to use for the new custom gate. Must
+            match the template gate type.
         :return: None
         """
         # First, check the gate type matches the template gate
@@ -72,6 +73,16 @@ def is_custom_gate(self, sample_id):
 
         return False
 
+    def remove_custom_gate(self, sample_id):
+        """
+        Removes a custom gate variation from the node. No error is thrown if a
+        custom gate for the sample ID does not exist.
+
+        :param sample_id: text string used to identify the custom gate
+        :return: None
+        """
+        self.custom_gates.pop(sample_id, None)
+
     def get_gate(self, sample_id=None):
         """
         Get Gate instance from GateNode. Specify sample_id to get sample custom gate.

src/flowkit/_models/gating_strategy.py

@@ -290,25 +290,16 @@ def rename_gate(self, gate_name, new_gate_name, gate_path=None):
         # rebuild DAG
         self._rebuild_dag()
 
-    def remove_gate(self, gate_name, gate_path=None, keep_children=False):
+    def _remove_template_gate(self, gate_node, keep_children=False):
         """
-        Remove a gate from the gating strategy. Any descendant gates will also be removed
-        unless keep_children=True. In all cases, if a BooleanGate exists that references
-        the gate to remove, a GateTreeError will be thrown indicating the BooleanGate
-        must be removed prior to removing the gate.
+        Handles case for removing template gate from gate tree.
 
-        :param gate_name: text string of a gate name
-        :param gate_path: complete ordered tuple of gate names for unique set of gate ancestors.
-            Required if gate_name is ambiguous
+        :param gate_node: GateNode to remove
         :param keep_children: Whether to keep child gates. If True, the child gates will be
             remapped to the removed gate's parent. Default is False, which will delete all
             descendant gates.
         :return: None
         """
-        # First, get the gate node from anytree
-        # Note, this will raise an error on ambiguous gates so no need
-        # to handle that case
-        gate_node = self._get_gate_node(gate_name, gate_path=gate_path)
         gate = gate_node.gate
 
         # single quadrants can't be removed, their "parent" QuadrantGate must be removed
@@ -327,7 +318,7 @@ def remove_gate(self, gate_name, gate_path=None, keep_children=False):
             s_gate = s_gate_node.gate
 
             if isinstance(s_gate, fk_gates.BooleanGate):
-                raise GateTreeError("BooleanGate %s references gate %s" % (s_gate.gate_name, gate_name))
+                raise GateTreeError("BooleanGate %s references gate %s" % (s_gate.gate_name, gate.gate_name))
 
         # At this point we're about to modify the tree and
         # removing a gate nullifies any previous results,
@@ -358,6 +349,37 @@ def remove_gate(self, gate_name, gate_path=None, keep_children=False):
         # Now, rebuild the DAG (easier than modifying it)
         self._rebuild_dag()
 
+    def remove_gate(self, gate_name, gate_path=None, sample_id=None, keep_children=False):
+        """
+        Remove a gate from the gating strategy. Any descendant gates will also be removed
+        unless keep_children=True. In all cases, if a BooleanGate exists that references
+        the gate to remove, a GateTreeError will be thrown indicating the BooleanGate
+        must be removed prior to removing the gate.
+
+        :param gate_name: text string of a gate name
+        :param gate_path: complete ordered tuple of gate names for unique set of gate ancestors.
+            Required if gate_name is ambiguous
+        :param sample_id: text string for Sample ID to remove only its custom Sample gate and
+            retain the template gate (and other custom gates if they exist).
+        :param keep_children: Whether to keep child gates. If True, the child gates will be
+            remapped to the removed gate's parent. Default is False, which will delete all
+            descendant gates.
+        :return: None
+        """
+        # First, get the gate node from anytree
+        # Note, this will raise an error on ambiguous gates so no need
+        # to handle that case
+        gate_node = self._get_gate_node(gate_name, gate_path=gate_path)
+
+        # determine whether user requested to remove the template gate or a custom Sample gate
+        if sample_id is None:
+            # Remove template gate, which removes the entire node from the tree
+            self._remove_template_gate(gate_node, keep_children=keep_children)
+        else:
+            # Remove custom sample gate, which is simpler.
+            # Stay silent if key doesn't exist.
+            gate_node.remove_custom_gate(sample_id)
+
     def add_transform(self, transform_id, transform):
         """
         Add a transform to the gating strategy, see `transforms` module. The transform ID must be unique in the
@@ -566,7 +588,7 @@ def get_max_depth(self):
     def get_gate_hierarchy(self, output='ascii', **kwargs):
         """
         Retrieve the hierarchy of gates in the gating strategy in several formats, including text,
-        dictionary, or JSON. If output == 'json', extra keyword arguments are passed to json.dumps
+        dictionary, or JSON. If output == 'json', extra keyword arguments are passed to 'json.dumps'
 
         :param output: Determines format of hierarchy returned, either 'ascii',
             'dict', or 'JSON' (default is 'ascii')

tests/gating_strategy_modify_gate_tree_tests.py

@@ -298,6 +298,37 @@ def test_remove_bool_dep(self):
             fk.exceptions.GateReferenceError, gs.get_gate, gate_name_to_remove
         )
 
+    def test_remove_custom_gate(self):
+        gs = copy.deepcopy(self.gating_strategy)
+
+        # make custom CD4 pos gate
+        custom_gate_name = "CD4-pos-poly"
+        gate_path = ('root', 'Time-range', 'Singlets-poly', 'Live-poly', 'CD3-pos-range')
+        new_poly_vertices = [[0.26, 0.37], [0.67, 0.6], [0.67, 0.8], [0.26, 0.8]]
+        cd4_gate_copy = copy.deepcopy(gs.get_gate(custom_gate_name))
+        cd4_gate_copy.vertices = new_poly_vertices
+
+        # add to gating strategy as custom gate for sample ID
+        sample_id = self.sample.id
+        gs.add_gate(cd4_gate_copy, gate_path, sample_id=sample_id)
+
+        # verify custom gate exists & vertices match
+        is_cd4_gate_custom = gs.is_custom_gate(sample_id, custom_gate_name)
+        self.assertTrue(is_cd4_gate_custom)
+        custom_gate_stored = gs.get_gate(custom_gate_name, sample_id=sample_id)
+        self.assertEqual(custom_gate_stored.vertices, new_poly_vertices)
+
+        # Remove just the custom gate
+        gs.remove_gate(custom_gate_name, gate_path, sample_id=sample_id)
+
+        # verify we can no longer access the custom gate for the sample
+        is_cd4_gate_custom = gs.is_custom_gate(sample_id, custom_gate_name)
+        self.assertFalse(is_cd4_gate_custom)
+
+        # Finally, verify the template gate still exists
+        template_gate = gs.get_gate(custom_gate_name)
+        self.assertIsInstance(template_gate, fk.gates.PolygonGate)
+
     def test_remove_gate_keep_children(self):
         # reminder of gate tree relevant to test:
         # root