Repository: cogent3/cogent3
Pull request #2080 (2024-10-30 08:00:43)

ENH: improve consistency of arrow head shape in feature drawing
## Summary by Sourcery

Enhance the feature drawing by improving the consistency of arrow head shapes through a new calculation method and add corresponding tests to validate the changes.

Enhancements:
- Improve consistency of arrow head shape in feature drawing by introducing a new function to calculate arrow width.

Tests:
- Add tests for the new arrow width calculation function to ensure correct behavior and error handling.


Patch:
diff --git a/src/cogent3/core/annotation.py b/src/cogent3/core/annotation.py
index af8f5aa4c..254546970 100644
--- a/src/cogent3/core/annotation.py
+++ b/src/cogent3/core/annotation.py
@@ -274,7 +274,7 @@ class Feature:
         """returns plotly trace"""
         from cogent3.draw.drawable import make_shape
 
-        return make_shape(type_=self)
+        return make_shape(type_=self, parent_length=len(self.parent))
 
     def to_dict(self):
         """returns"""
diff --git a/src/cogent3/draw/drawable.py b/src/cogent3/draw/drawable.py
index a2fbbd70a..8416a63b1 100644
--- a/src/cogent3/draw/drawable.py
+++ b/src/cogent3/draw/drawable.py
@@ -1,5 +1,6 @@
 import os
 import pathlib
+import typing
 
 import numpy
 
@@ -554,6 +555,7 @@ class Shape:
         showlegend=True,
         hoverinfo=None,
         fillcolor=None,
+        **kwargs,
     ):
         self.filled = filled
         self.fillcolor = fillcolor
@@ -672,58 +674,126 @@ class Diamond(Shape):
         self.y = numpy.array(ys)
 
 
+def _calc_arrow_width(
+    *,
+    parent_length: int,
+    feature_width: float,
+    frac: float = 0.05,
+) -> float:
+    aw = parent_length * frac
+    if aw > feature_width:
+        aw = feature_width * frac
+    return aw
+
+
+def _make_rectangles(coords: list[list[int]], y=0, height=0.25) -> tuple[list[float]]:
+    xs = []
+    ys = []
+    y_coord = [y, y + height, y + height, y, y]
+    for coord in coords:
+        # Add coordinates for individual rectangle
+        width = abs(coord[0] - coord[1])
+        x_coord = min(coord[0], coord[1])
+        xs.append([x_coord, x_coord, x_coord + width, x_coord + width, x_coord])
+        ys.append(y_coord)
+    return xs, ys
+
+
+def _make_arrow_head(
+    coord: list[float],
+    y: float,
+    height: float,
+    head_width_frac: float,
+    reverse: bool,
+    parent_length: int,
+) -> tuple[list[float]]:
+    width = abs(coord[0] - coord[1])
+    x_coord = min(coord[0], coord[1])
+    hh = height * head_width_frac * 2
+    hw = _calc_arrow_width(
+        parent_length=parent_length, feature_width=width, frac=head_width_frac
+    )
+
+    # Coordinates for arrow head
+    arrow_x = [
+        x_coord,
+        x_coord + width - hw,
+        x_coord + width - hw,
+        x_coord + width,
+        x_coord + width - hw,
+        x_coord + width - hw,
+        x_coord,
+        x_coord,
+    ]
+    arrow_y = [
+        y,
+        y,
+        y - hh,
+        y + height / 2,
+        y + height + hh,
+        y + height,
+        y + height,
+        y,
+    ]
+    if reverse:
+        arrow_x = numpy.array(arrow_x)
+        arrow_x = list(numpy.flip(arrow_x.max() - arrow_x + arrow_x.min()))
+        arrow_y = list(numpy.flip(arrow_y))
+    return arrow_x, arrow_y
+
+
+def _connecting_lines(
+    x_grp: list[list[float]], y_grp: list[list[float]], connect_y: float
+) -> tuple[list[typing.Union[float, None]]]:
+    new_xs, new_ys = x_grp[0][:], y_grp[0][:]
+    y_connect_coord = [None, connect_y, connect_y, None]
+    for i, (xs, ys) in enumerate(zip(x_grp[1:], y_grp[1:]), 1):
+        new_xs.extend([None, max(x_grp[i - 1]), min(x_grp[i]), None] + xs)
+        new_ys.extend(y_connect_coord + ys)
+    return new_xs, new_ys
+
+
 class Arrow(Shape):
     def __init__(
-        self, coords, y=0, height=0.25, arrow_head_w=0.1, reverse=False, **kwargs
+        self,
+        coords,
+        y=0,
+        height=0.25,
+        arrow_head_w=0.05,
+        reverse=False,
+        parent_length: typing.Optional[int] = None,
+        **kwargs,
     ):
-        super(Arrow, self).__init__(**kwargs)
-        xs = []
-        ys = []
-        for i in range(len(coords) - 1):
-            # Add coordinates for individual rectangle
-            width = abs(coords[i][0] - coords[i][1])
-            x_coord = min(coords[i][0], coords[i][1])
-            xs += [x_coord, x_coord, x_coord + width, x_coord + width, x_coord]
-            ys += [y, y + height, y + height, y, y]
-            # Add coordinates for connecting line segment
-            xs += [None, coords[i][1], coords[i + 1][0], None]
-            ys += [None, y + height / 2, y + height / 2, None]
+        super().__init__(**kwargs)
 
-        width = abs(coords[-1][0] - coords[-1][1])
-        x_coord = min(coords[-1][0], coords[-1][1])
-        hh = height * arrow_head_w * 2
-        hw = width * arrow_head_w * 2
+        if len(coords) == 1 or reverse:
+            head_index = 0
+            rec_coords = None if len(coords) == 1 else coords[1:]
+        else:
+            head_index = -1
+            rec_coords = coords[:-1]
 
-        # Coordinates for arrow head
-        arrow_x = [
-            x_coord,
-            x_coord + width - hw,
-            x_coord + width - hw,
-            x_coord + width,
-            x_coord + width - hw,
-            x_coord + width - hw,
-            x_coord,
-            x_coord,
-        ]
-        arrow_y = [
-            y,
-            y,
-            y - hh,
-            y + height / 2,
-            y + height + hh,
-            y + height,
-            y + height,
-            y,
-        ]
-        if not reverse:
-            xs += arrow_x
-            ys += arrow_y
+        arrow_x, arrow_y = _make_arrow_head(
+            coords[head_index],
+            y=y,
+            height=height,
+            head_width_frac=arrow_head_w,
+            reverse=reverse,
+            parent_length=parent_length,
+        )
+        if rec_coords:
+            xs, ys = _make_rectangles(rec_coords, y=y, height=height)
+        else:
+            xs, ys = [], []
+
+        if reverse:
+            xs.insert(0, arrow_x)
+            ys.insert(0, arrow_y)
         else:
-            arrow_x = numpy.array(arrow_x)
-            arrow_y = numpy.array(arrow_y)
-            xs += list(numpy.flip(arrow_x.max() - arrow_x + arrow_x.min()))
-            ys += list(numpy.flip(arrow_y))
+            xs.append(arrow_x)
+            ys.append(arrow_y)
 
+        xs, ys = _connecting_lines(xs, ys, y + height / 2)
         self.x = numpy.array(xs)
         self.y = numpy.array(ys)
 


Test patch:
diff --git a/tests/test_draw/test_draw_integration.py b/tests/test_draw/test_draw_integration.py
index b2129376d..ab3bf7e22 100644
--- a/tests/test_draw/test_draw_integration.py
+++ b/tests/test_draw/test_draw_integration.py
@@ -5,7 +5,12 @@ from numpy.testing import assert_allclose
 
 from cogent3 import load_aligned_seqs, make_aligned_seqs, make_table
 from cogent3.core.annotation_db import GffAnnotationDb
-from cogent3.draw.drawable import AnnotatedDrawable, Drawable, get_domain
+from cogent3.draw.drawable import (
+    AnnotatedDrawable,
+    Drawable,
+    _calc_arrow_width,
+    get_domain,
+)
 from cogent3.util.union_dict import UnionDict
 
 
@@ -391,7 +396,8 @@ class AlignmentDrawablesTests(BaseDrawablesTests):
         db.add_feature(seqid="b", biotype="gene", name="1", spans=[(1, 5)])
         db.add_feature(seqid="b", biotype="gene", name="1", spans=[(5, 1)])
         aln.annotation_db = db
-        aln.get_drawable()
+        g = aln.get_drawable()
+        assert isinstance(g, Drawable)
 
 
 class TableDrawablesTest(BaseDrawablesTests):
@@ -405,5 +411,8 @@ class TableDrawablesTest(BaseDrawablesTests):
         self._check_drawable_attrs(drawable.figure, "table")
 
 
-if __name__ == "__main__":
-    unittest.main()
+def test_calculating_arrow_width_adjusted():
+    aw = _calc_arrow_width(parent_length=100, feature_width=10, frac=1.0)
+    assert aw == 10
+    aw = _calc_arrow_width(parent_length=100, feature_width=10, frac=0.05)
+    assert aw == 5