From 2b582fc2a01165ad505c6bc4ddac54fed7cd574f Mon Sep 17 00:00:00 2001
From: mbsantiago <santiago.mbal@gmail.com>
Date: Sat, 8 Aug 2026 13:39:32 +0100
Subject: [PATCH] Fix typing issues

---
 src/batdetect2/outputs/formats/parquet.py | 9 +++++++--
 src/batdetect2/plotting/detections.py     | 8 ++++----
 src/batdetect2/plotting/heatmaps.py       | 4 ++--
 3 files changed, 13 insertions(+), 8 deletions(-)

diff --git a/src/batdetect2/outputs/formats/parquet.py b/src/batdetect2/outputs/formats/parquet.py
index 909fa9d..fdd3e7b 100644
--- a/src/batdetect2/outputs/formats/parquet.py
+++ b/src/batdetect2/outputs/formats/parquet.py
@@ -1,5 +1,5 @@
 from pathlib import Path
-from typing import List, Literal, Sequence
+from typing import List, Literal, Sequence, TypedDict
 from uuid import UUID
 
 import numpy as np
@@ -26,6 +26,11 @@ class ParquetOutputConfig(BaseConfig):
     include_geometry: bool = True
 
 
+class ClipInfo(TypedDict):
+    clip: data.Clip
+    preds: list[Detection]
+
+
 class ParquetFormatter(OutputFormatterProtocol[ClipDetections]):
     def __init__(
         self,
@@ -120,7 +125,7 @@ class ParquetFormatter(OutputFormatterProtocol[ClipDetections]):
         else:
             df = pd.read_parquet(path)
 
-        predictions_by_clip = {}
+        predictions_by_clip: dict[UUID, ClipInfo] = {}
 
         for _, row in df.iterrows():
             clip_uuid = row["clip_uuid"]
diff --git a/src/batdetect2/plotting/detections.py b/src/batdetect2/plotting/detections.py
index f25a111..ef96532 100644
--- a/src/batdetect2/plotting/detections.py
+++ b/src/batdetect2/plotting/detections.py
@@ -185,25 +185,25 @@ def plot_clip_evaluation(
                     label="found GT",
                     edgecolor=gt_color,
                     facecolor="none" if not fill else gt_color,
-                    linestyle=gt_linestyle,
+                    linestyle=gt_linestyle,  # type: ignore
                 ),
                 patches.Patch(
                     label="missed GT",
                     edgecolor=missed_gt_color,
                     facecolor="none" if not fill else missed_gt_color,
-                    linestyle=missed_gt_linestyle,
+                    linestyle=missed_gt_linestyle,  # type: ignore
                 ),
                 patches.Patch(
                     label="true Det",
                     edgecolor=true_pred_color,
                     facecolor="none" if not fill else true_pred_color,
-                    linestyle=true_pred_linestyle,
+                    linestyle=true_pred_linestyle,  # type: ignore
                 ),
                 patches.Patch(
                     label="false Det",
                     edgecolor=false_pred_color,
                     facecolor="none" if not fill else false_pred_color,
-                    linestyle=false_pred_linestyle,
+                    linestyle=false_pred_linestyle,  # type: ignore
                 ),
             ]
         )
diff --git a/src/batdetect2/plotting/heatmaps.py b/src/batdetect2/plotting/heatmaps.py
index 9f4021f..8f10330 100644
--- a/src/batdetect2/plotting/heatmaps.py
+++ b/src/batdetect2/plotting/heatmaps.py
@@ -1,9 +1,9 @@
 """Plot heatmaps."""
 
+import matplotlib.pyplot as plt
 import numpy as np
 import torch
 from matplotlib import axes, patches
-from matplotlib.cm import get_cmap
 from matplotlib.colors import Colormap, LinearSegmentedColormap, to_rgba
 
 from batdetect2.plotting.common import create_ax
@@ -80,7 +80,7 @@ def plot_classification_heatmap(
         raise ValueError("Inconsistent number of class names")
 
     if not isinstance(cmap, Colormap):
-        cmap = get_cmap(cmap)
+        cmap = plt.get_cmap(cmap)
 
     handles = []