Fix unifurcations when a 0-length leaf is removed (#108)

neuror/sanitize.py

@@ -49,8 +49,6 @@ def sanitize(input_neuron, output_path):
     if neuron.soma.type == SomaType.SOMA_UNDEFINED:  # pylint: disable=no-member
         raise CorruptedMorphology(f'{input_neuron} has an invalid or no soma.')
 
-    neuron.remove_unifurcations()
-
     for section in neuron.iter():
         section.diameters = np.clip(section.diameters, 0, None)
 
@@ -64,11 +62,14 @@ def sanitize(input_neuron, output_path):
                                           f'{section.parent.id}) (type: {section.parent.type})')
 
     try:
-        fix_non_zero_segments(neuron).write(str(output_path))
+        sanitized_neuron = fix_non_zero_segments(neuron)
     except ZeroLengthRootSection as e:
         # reraise to attach the morphology path
         raise ZeroLengthRootSection(f"Failed morphology: {input_neuron}") from e
 
+    sanitized_neuron.remove_unifurcations()
+    sanitized_neuron.write(str(output_path))
+
 
 def _sanitize_one(path, input_folder, output_folder):
     '''Function to be called by sanitize_all to catch all exceptions

tests/test_sanitize.py

@@ -201,6 +201,32 @@ def test_sanitize__negative_diameters():
         )
 
 
+def test_sanitize__zero_length_leaf():
+
+    content = (
+    """
+     1 1  0  0 0 1. -1
+     2 2  0  0 0 1.  1
+     3 2  1  0 0 1.  2
+     4 2  2  0 0 1.  3
+     5 2  1  0 0 1.  3  # 0-length leaf
+    """
+    )
+    with _tmp_file(content, extension="swc") as input_filepath, \
+         _tmp_file("", extension="swc") as output_filepath:
+
+        tested.sanitize(input_filepath, output_filepath)
+
+        assert_array_equal(
+            Morphology(output_filepath).points,
+            [
+                [0, 0, 0],
+                [1, 0, 0],
+                [2, 0, 0],
+                # Here the leaf was completly deleted
+            ],
+        )
+
 
 def test_sanitize_all(tmpdir):
     tmpdir = Path(tmpdir)