chore: fix typing issues

zapata-engineering · Oct 19, 2023 · 31a6dc3 · 31a6dc3
1 parent 7888b0b
commit 31a6dc3
Show file tree

Hide file tree

Showing 2 changed files with 8 additions and 3 deletions.
diff --git 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
@@ -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
     )