From e282e6f6e67067a4d0d1a224435f732cb2eb9e10 Mon Sep 17 00:00:00 2001
From: zm711 <92116279+zm711@users.noreply.github.com>
Date: Tue, 22 Aug 2023 17:17:22 -0400
Subject: [PATCH] add typing to more functions

---
 src/spikeanalysis/intrinsic_plotter.py |  9 +++++++--
 src/spikeanalysis/spike_analysis.py    |  2 +-
 src/spikeanalysis/stimulus_data.py     | 15 +++++++++++++--
 3 files changed, 21 insertions(+), 5 deletions(-)

diff --git a/src/spikeanalysis/intrinsic_plotter.py b/src/spikeanalysis/intrinsic_plotter.py
index cc4951f..760f2a0 100644
--- a/src/spikeanalysis/intrinsic_plotter.py
+++ b/src/spikeanalysis/intrinsic_plotter.py
@@ -324,7 +324,7 @@ def _plot_cdf_pdf(self, df: pd.DataFrame):
 
     def _generate_amp_depth_bins(
         self,
-        sp,
+        sp: SpikeData,
         spike_amps: np.ndarray,
         probe_len: float,
         pitch: float,
@@ -349,7 +349,12 @@ def _generate_amp_depth_bins(
         return depth_bins, amp_bins, recording_duration
 
     def _compute_cdf_pdf(
-        self, spike_amps, spike_depths, amp_bins, depth_bins, recording_dur
+        self,
+        spike_amps: np.ndarray,
+        spike_depths: np.ndarray,
+        amp_bins: np.ndarray,
+        depth_bins: np.ndarray,
+        recording_dur: float,
     ) -> tuple[np.ndarray, np.ndarray]:
         n_depth_bins = len(depth_bins) - 1
         n_amp_bins = len(amp_bins) - 1
diff --git a/src/spikeanalysis/spike_analysis.py b/src/spikeanalysis/spike_analysis.py
index aa03fec..a73f7b0 100644
--- a/src/spikeanalysis/spike_analysis.py
+++ b/src/spikeanalysis/spike_analysis.py
@@ -869,7 +869,7 @@ def get_responsive_neurons(self, z_parameters: Optional[dict] = None):
 
                 self.responsive_neurons[stim][key] = responsive_neurons
 
-    def _merge_events(self, event_0, event_1):
+    def _merge_events(self, event_0: dict, event_1: dict):
         """Utility function for merging digital and analog events into one dictionary"""
         events = {**event_0, **event_1}
         return events
diff --git a/src/spikeanalysis/stimulus_data.py b/src/spikeanalysis/stimulus_data.py
index c81114e..88853db 100644
--- a/src/spikeanalysis/stimulus_data.py
+++ b/src/spikeanalysis/stimulus_data.py
@@ -131,6 +131,17 @@ def create_neo_reader(self, file_name: Optional[str] = None):
         """
         Function that creates a Neo IntanRawIO reader and then parses the header
 
+        Parameters
+        ----------
+        file_name: Optional[filename]
+            Default None uses the internal function, whereas providing a file_name will call
+            the neo function `get_rawio_class` to return a neo reader that can be used
+
+        Returns
+        -------
+        neo_class: neo.rawio
+            Returns a neo.rawio class if filename is given otherwise stores values internally
+
         """
         if file_name is None:
             reader = neo.rawio.IntanRawIO(filename=self._filename)
@@ -146,7 +157,7 @@ def create_neo_reader(self, file_name: Optional[str] = None):
         self.sample_frequency = sample_freq
         self.reader = reader
 
-    def get_analog_data(self, time_slice=(None, None)):
+    def get_analog_data(self, time_slice: tuple = (None, None)):
         """
         Function to load analog data from an Intan file. Requires the IntanRawIO to be generated with
         `create_neo_reader`
@@ -259,7 +270,7 @@ def _valueround(self, x: float, precision: int = 2, base: float = 0.25) -> float
         """
         return round(base * round(float(x) / base), precision)
 
-    def get_raw_digital_data(self, time_slice=(None, None)):
+    def get_raw_digital_data(self, time_slice: tuple = (None, None)):
         """
         This is a function that in the future will get the digital data, but currently due
         to the inability to grab digital from neo automatically. Calls on internal hack to