From 31a6dc38d04a120de0e242fe6ef4149bebf7d9e3 Mon Sep 17 00:00:00 2001
From: James Clark <james.clark@zapatacomputing.com>
Date: Thu, 19 Oct 2023 14:22:40 +0100
Subject: [PATCH] chore: fix typing issues

---
 src/orqviz/plot_utils.py  | 9 ++++++---
 src/orqviz/scans/plots.py | 2 ++
 2 files changed, 8 insertions(+), 3 deletions(-)

diff --git a/src/orqviz/plot_utils.py b/src/orqviz/plot_utils.py
index c9190d3..0643264 100644
--- a/src/orqviz/plot_utils.py
+++ b/src/orqviz/plot_utils.py
@@ -4,6 +4,7 @@
 import matplotlib.pyplot as plt
 import matplotlib.ticker as tck
 import numpy as np
+from matplotlib.cm import ScalarMappable
 
 
 def normalize_color_and_colorbar(
@@ -40,9 +41,11 @@ def normalize_color_and_colorbar(
                 "Provided ax does not contain an image in ax.images or ax.collections"
             ) from e
 
-    image.colorbar.remove()
-    image.set_clim(vmin=min_val, vmax=max_val)
-    image.set_cmap(cmap)
+    assert isinstance(image, ScalarMappable)
+    if image.colorbar is not None:
+        image.colorbar.remove()
+    image.set_clim(vmin=min_val, vmax=max_val)  # type: ignore
+    image.set_cmap(cmap)  # type: ignore
     fig.colorbar(image, ax=ax)
 
 
diff --git a/src/orqviz/scans/plots.py b/src/orqviz/scans/plots.py
index 444a679..368ef3c 100644
--- a/src/orqviz/scans/plots.py
+++ b/src/orqviz/scans/plots.py
@@ -163,6 +163,8 @@ def plot_2D_scan_result_as_3D(
     """
     ax = _check_and_create_3D_ax(ax=ax)
 
+    assert ax is not None
+
     x, y = scan2D_result._get_coordinates_on_directions(
         in_units_of_direction=in_units_of_direction
     )