From eda0f30ca8bcafb4066698e0c927c5b8f5252772 Mon Sep 17 00:00:00 2001
From: Leonardo Schwarz <leonardo.schwarz@fgcz.ethz.ch>
Date: Mon, 28 Oct 2024 09:26:08 +0100
Subject: [PATCH] fix some type hints

---
 src/depiction/persistence/imzml/imzml_writer.py | 13 ++++++-------
 src/depiction/persistence/types.py              |  4 ++--
 2 files changed, 8 insertions(+), 9 deletions(-)

diff --git a/src/depiction/persistence/imzml/imzml_writer.py b/src/depiction/persistence/imzml/imzml_writer.py
index 3747ec6..c521c8b 100644
--- a/src/depiction/persistence/imzml/imzml_writer.py
+++ b/src/depiction/persistence/imzml/imzml_writer.py
@@ -1,17 +1,16 @@
 from __future__ import annotations
 
-from pathlib import Path
-from typing import TYPE_CHECKING
-
 import pyimzml
 import pyimzml.ImzMLWriter
+from pathlib import Path
+from typing import TYPE_CHECKING
 
 from depiction.persistence.imzml.imzml_alignment_tracker import ImzmlAlignmentTracker
 from depiction.persistence.imzml.imzml_mode_enum import ImzmlModeEnum
 from depiction.persistence.types import GenericWriter
 
 if TYPE_CHECKING:
-    import numpy as np
+    from numpy.typing import NDArray
 
 
 class ImzmlWriter(GenericWriter):
@@ -65,9 +64,9 @@ def is_aligned(self) -> bool:
 
     def add_spectrum(
         self,
-        mz_arr: np.ndarray,
-        int_arr: np.ndarray,
-        coordinates: tuple[int, int] | tuple[int, int, int],
+        mz_arr: NDArray[float],
+        int_arr: NDArray[float],
+        coordinates: tuple[int, int] | tuple[int, int, int] | NDArray[int],
     ) -> None:
         if len(mz_arr) != len(int_arr):
             raise ValueError(f"{len(mz_arr)=} and {len(int_arr)=} must be equal.")
diff --git a/src/depiction/persistence/types.py b/src/depiction/persistence/types.py
index 9e627d6..dd897eb 100644
--- a/src/depiction/persistence/types.py
+++ b/src/depiction/persistence/types.py
@@ -67,7 +67,7 @@ def get_spectrum(self, i_spectrum: int) -> tuple[NDArray[float], NDArray[float]]
         """Returns the m/z and intensity arrays of the i-th spectrum."""
         return self.get_spectrum_mz(i_spectrum=i_spectrum), self.get_spectrum_int(i_spectrum=i_spectrum)
 
-    def get_spectrum_with_coords(self, i_spectrum: int) -> tuple[NDArray[float], NDArray[float], NDArray[float]]:
+    def get_spectrum_with_coords(self, i_spectrum: int) -> tuple[NDArray[float], NDArray[float], NDArray[int]]:
         """Returns the m/z, intensity and v arrays of the i-th spectrum."""
         mz_arr = self.get_spectrum_mz(i_spectrum=i_spectrum)
         int_arr = self.get_spectrum_int(i_spectrum=i_spectrum)
@@ -215,7 +215,7 @@ def add_spectrum(
         self,
         mz_arr: np.ndarray,
         int_arr: np.ndarray,
-        coordinates: tuple[int, int] | tuple[int, int, int],
+        coordinates: tuple[int, int] | tuple[int, int, int] | NDArray[int],
     ) -> None: ...
 
     def copy_spectra(