refine(key_points): address PR #2343 review feedback

- Generalize TRACK error msg in _resolve_color_idx to not reference
  "Detections object" (function now shared with keypoint annotators)
- Consolidate multi-line error strings to full line width
- Add Raises section to _resolve_color_idx docstring
- Add TRACK-not-supported note to VertexAnnotator, EdgeAnnotator,
  VertexLabelAnnotator color_lookup docs
- Add ColorPalette+ColorLookup tests: INDEX/CLASS/KEYPOINT for all
  three keypoint annotators + CLASS-without-class_id error case
- Fix _resolve_color_idx signature: NDArray[np.generic] (linter)

[resolve items 1,2,3,5,6] PR #2343 review by @Copilot and @Borda
Challenge: evidence=VALID suggestion=VALID resolution=as-suggested

---
Co-authored-by: claude[bot] <209825114+claude[bot]@users.noreply.github.com>

Author: Borda

src/supervision/annotators/utils.py

@@ -3,7 +3,7 @@
 import re
 import textwrap
 from enum import Enum
-from typing import Any
+from typing import Any, cast
 
 import numpy as np
 import numpy.typing as npt
@@ -44,8 +44,8 @@ def _resolve_color_idx(
     instance_idx: int,
     color_lookup: ColorLookup,
     count: int,
-    class_id: npt.NDArray[np.int_] | None = None,
-    tracker_id: npt.NDArray[np.int_] | None = None,
+    class_id: npt.NDArray[np.generic] | None = None,
+    tracker_id: npt.NDArray[np.generic] | None = None,
     keypoint_idx: int | None = None,
 ) -> int:
     """Resolve a palette index from raw field arrays.
@@ -63,6 +63,14 @@ def _resolve_color_idx(
 
     Returns:
         An integer index suitable for ``ColorPalette.by_idx()``.
+
+    Raises:
+        ValueError: If ``instance_idx`` is out of bounds for the given ``count``.
+        ValueError: If ``color_lookup`` is ``CLASS`` and ``class_id`` is ``None``.
+        ValueError: If ``color_lookup`` is ``TRACK`` and ``tracker_id`` is ``None``.
+        ValueError: If ``color_lookup`` is ``KEYPOINT`` and ``keypoint_idx`` is
+            ``None``.
+        ValueError: If ``color_lookup`` is an unsupported strategy.
     """
     if instance_idx >= count:
         raise ValueError(
@@ -74,17 +82,15 @@ def _resolve_color_idx(
     elif color_lookup == ColorLookup.CLASS:
         if class_id is None:
             raise ValueError(
-                "Could not resolve color by class because "
-                "class_id is not available. Try setting "
-                "color_lookup to sv.ColorLookup.INDEX."
+                "Could not resolve color by class because class_id is not available. "
+                "Try setting color_lookup to sv.ColorLookup.INDEX."
             )
         return int(class_id[instance_idx])
     elif color_lookup == ColorLookup.TRACK:
         if tracker_id is None:
             raise ValueError(
-                "Could not resolve color by track because "
-                "tracker_id is not available. Make sure that the "
-                "Detections object contains tracker_id data."
+                "Could not resolve color by track because tracker_id is not available. "
+                "Make sure tracker_id is set on the input object."
             )
         return int(tracker_id[instance_idx])
     elif color_lookup == ColorLookup.KEYPOINT:
@@ -434,7 +440,7 @@ def snap_boxes(
     bottom_shift = height - result[bottom_overflow, 3]
     result[bottom_overflow, 1:4:2] += bottom_shift[:, np.newaxis]
 
-    return result.astype(np.float32)  # type: ignore
+    return cast(np.ndarray[Any, np.dtype[np.float32]], result.astype(np.float32))
 
 
 class Trace:

src/supervision/key_points/annotators.py

@@ -78,7 +78,8 @@ def __init__(
             color_lookup: Strategy for mapping colors to annotations.
                 Options are `INDEX` (per-skeleton index), `CLASS`
                 (per class_id), and `KEYPOINT` (per keypoint index within
-                each skeleton).
+                each skeleton). Note: ``TRACK`` is not supported for
+                keypoint annotators.
         """
         self.color = color
         self.radius = radius
@@ -181,7 +182,8 @@ def __init__(
             color_lookup: Strategy for mapping colors to annotations.
                 Options are `INDEX` (per-skeleton index), `CLASS`
                 (per class_id), and `KEYPOINT` (per keypoint index —
-                edge inherits the color of its first endpoint).
+                edge inherits the color of its first endpoint). Note:
+                ``TRACK`` is not supported for keypoint annotators.
         """
         self.color = color
         self.thickness = thickness
@@ -797,7 +799,8 @@ def __init__(
             color_lookup: Strategy for mapping colors to annotations.
                 Options are `INDEX` (per-skeleton index), `CLASS`
                 (per class_id), and `KEYPOINT` (per keypoint index within
-                each skeleton).
+                each skeleton). Note: ``TRACK`` is not supported for
+                keypoint annotators.
         """
         if isinstance(color, list):
             warn_deprecated(

tests/key_points/test_annotators.py

@@ -467,3 +467,164 @@ def test_resolve_labels_returns_expected(
     def test_resolve_labels_raises(self, labels, points_count, class_id, match):
         with pytest.raises(ValueError, match=match):
             sv.VertexLabelAnnotator._resolve_labels(labels, points_count, class_id)
+
+
+class TestVertexAnnotatorColorLookup:
+    """Verify VertexAnnotator respects each ColorLookup strategy with a ColorPalette."""
+
+    @pytest.fixture
+    def key_points_with_class(self) -> sv.KeyPoints:
+        """Two-instance, three-keypoint set with class_id set."""
+        return sv.KeyPoints(
+            xy=np.array(
+                [
+                    [[20.0, 20.0], [40.0, 40.0], [60.0, 60.0]],
+                    [[25.0, 25.0], [45.0, 45.0], [65.0, 65.0]],
+                ],
+                dtype=np.float32,
+            ),
+            class_id=np.array([0, 1], dtype=int),
+        )
+
+    @pytest.mark.parametrize(
+        "color_lookup",
+        [
+            pytest.param(sv.ColorLookup.INDEX, id="index"),
+            pytest.param(sv.ColorLookup.CLASS, id="class"),
+            pytest.param(sv.ColorLookup.KEYPOINT, id="keypoint"),
+        ],
+    )
+    def test_annotate_with_color_palette_returns_ndarray(
+        self, scene, key_points_with_class, color_lookup
+    ):
+        """ColorPalette + each ColorLookup produces a modified ndarray output."""
+        annotator = sv.VertexAnnotator(
+            color=sv.ColorPalette.DEFAULT,
+            radius=5,
+            color_lookup=color_lookup,
+        )
+        result = annotator.annotate(
+            scene=scene.copy(), key_points=key_points_with_class
+        )
+
+        assert isinstance(result, np.ndarray)
+        assert result.shape == scene.shape
+        assert not np.array_equal(result, scene)
+
+    def test_annotate_class_lookup_raises_when_class_id_none(self, scene):
+        """CLASS strategy raises ValueError when key_points.class_id is None."""
+        key_points = sv.KeyPoints(
+            xy=np.array([[[30.0, 30.0], [50.0, 50.0]]], dtype=np.float32),
+        )
+        annotator = sv.VertexAnnotator(
+            color=sv.ColorPalette.DEFAULT,
+            color_lookup=sv.ColorLookup.CLASS,
+        )
+
+        with pytest.raises(ValueError, match="class_id"):
+            annotator.annotate(scene=scene.copy(), key_points=key_points)
+
+
+class TestEdgeAnnotatorColorLookup:
+    """Verify EdgeAnnotator respects each ColorLookup strategy with a ColorPalette."""
+
+    @pytest.fixture
+    def key_points_triangle(self) -> sv.KeyPoints:
+        """Single-instance, three-vertex triangle useful with explicit edges."""
+        return sv.KeyPoints(
+            xy=np.array(
+                [[[10.0, 10.0], [80.0, 10.0], [45.0, 80.0]]],
+                dtype=np.float32,
+            ),
+            class_id=np.array([0], dtype=int),
+        )
+
+    @pytest.mark.parametrize(
+        "color_lookup",
+        [
+            pytest.param(sv.ColorLookup.INDEX, id="index"),
+            pytest.param(sv.ColorLookup.CLASS, id="class"),
+            pytest.param(sv.ColorLookup.KEYPOINT, id="keypoint"),
+        ],
+    )
+    def test_annotate_with_color_palette_returns_ndarray(
+        self, scene, key_points_triangle, color_lookup
+    ):
+        """ColorPalette + each ColorLookup produces a modified ndarray output."""
+        annotator = sv.EdgeAnnotator(
+            color=sv.ColorPalette.DEFAULT,
+            thickness=2,
+            edges=[(1, 2), (2, 3), (1, 3)],
+            color_lookup=color_lookup,
+        )
+        result = annotator.annotate(scene=scene.copy(), key_points=key_points_triangle)
+
+        assert isinstance(result, np.ndarray)
+        assert result.shape == scene.shape
+        assert not np.array_equal(result, scene)
+
+    def test_annotate_class_lookup_raises_when_class_id_none(self, scene):
+        """CLASS strategy raises ValueError when key_points.class_id is None."""
+        key_points = sv.KeyPoints(
+            xy=np.array([[[10.0, 10.0], [80.0, 10.0]]], dtype=np.float32),
+        )
+        annotator = sv.EdgeAnnotator(
+            color=sv.ColorPalette.DEFAULT,
+            edges=[(1, 2)],
+            color_lookup=sv.ColorLookup.CLASS,
+        )
+
+        with pytest.raises(ValueError, match="class_id"):
+            annotator.annotate(scene=scene.copy(), key_points=key_points)
+
+
+class TestVertexLabelAnnotatorColorLookup:
+    """Verify VertexLabelAnnotator respects each ColorLookup strategy."""
+
+    @pytest.fixture
+    def key_points_with_class(self) -> sv.KeyPoints:
+        """Two-instance, two-keypoint set with class_id set."""
+        return sv.KeyPoints(
+            xy=np.array(
+                [[[20.0, 20.0], [60.0, 60.0]], [[25.0, 25.0], [65.0, 65.0]]],
+                dtype=np.float32,
+            ),
+            class_id=np.array([0, 1], dtype=int),
+        )
+
+    @pytest.mark.parametrize(
+        "color_lookup",
+        [
+            pytest.param(sv.ColorLookup.INDEX, id="index"),
+            pytest.param(sv.ColorLookup.CLASS, id="class"),
+            pytest.param(sv.ColorLookup.KEYPOINT, id="keypoint"),
+        ],
+    )
+    def test_annotate_with_color_palette_returns_ndarray(
+        self, scene, key_points_with_class, color_lookup
+    ):
+        """ColorPalette + each ColorLookup produces a modified ndarray output."""
+        annotator = sv.VertexLabelAnnotator(
+            color=sv.ColorPalette.DEFAULT,
+            color_lookup=color_lookup,
+        )
+        result = annotator.annotate(
+            scene=scene.copy(), key_points=key_points_with_class
+        )
+
+        assert isinstance(result, np.ndarray)
+        assert result.shape == scene.shape
+        assert not np.array_equal(result, scene)
+
+    def test_annotate_class_lookup_raises_when_class_id_none(self, scene):
+        """CLASS strategy raises ValueError when key_points.class_id is None."""
+        key_points = sv.KeyPoints(
+            xy=np.array([[[30.0, 30.0], [50.0, 50.0]]], dtype=np.float32),
+        )
+        annotator = sv.VertexLabelAnnotator(
+            color=sv.ColorPalette.DEFAULT,
+            color_lookup=sv.ColorLookup.CLASS,
+        )
+
+        with pytest.raises(ValueError, match="class_id"):
+            annotator.annotate(scene=scene.copy(), key_points=key_points)