From 9be188312599c7ce2e847016df8380aa38f2105e Mon Sep 17 00:00:00 2001
From: zm711 <92116279+zm711@users.noreply.github.com>
Date: Thu, 11 Jan 2024 14:49:34 -0500
Subject: [PATCH] add return option

---
 src/spikeanalysis/spike_analysis.py | 10 ++++++++++
 test/test_spike_analysis.py         | 18 ++++++++++++++++++
 2 files changed, 28 insertions(+)

diff --git a/src/spikeanalysis/spike_analysis.py b/src/spikeanalysis/spike_analysis.py
index 3bce7ed..711f8d6 100644
--- a/src/spikeanalysis/spike_analysis.py
+++ b/src/spikeanalysis/spike_analysis.py
@@ -819,6 +819,16 @@ def autocorrelogram(self, time_ms: float = 500):
 
         self.acg = acg
 
+    def return_value(self, value: str):
+
+        _values = ("zscores", "raw_zscores", "mean_firing_rate", "raw_firing_rate", "correlations", "latency", "psths")
+
+        if hasattr(self, value):
+            return getattr(self, value)
+        else:
+            print(f"possible values are {_values}")
+            raise AttributeError(f"{value} does not exist run appropriate function")
+
     def _generate_sample_z_parameter(self, save: bool = True) -> dict:
         """
         Function for providing example z score parameters. Then saves as json
diff --git a/test/test_spike_analysis.py b/test/test_spike_analysis.py
index 4e15fb1..635aded 100644
--- a/test/test_spike_analysis.py
+++ b/test/test_spike_analysis.py
@@ -164,6 +164,24 @@ def test_z_score_data(sa):
     assert np.sum(z_data[0, 0, :15]) < np.sum(z_data[0, 0, 150:200]), "Should be high z score"
 
 
+def test_return_value(sa):
+    sa.events = {
+        "0": {
+            "events": np.array([100, 200]),
+            "lengths": np.array([100, 100]),
+            "trial_groups": np.array([1, 1]),
+            "stim": "test",
+        }
+    }
+    sa.get_raw_psth(window=[0, 300], time_bin_ms=50)
+
+    psths = sa.return_value("psths")
+    assert "test" in psths.keys()
+
+    with pytest.raises(AttributeError):
+        _ = sa.return_value("lats")
+
+
 def test_get_interspike_intervals(sa):
     sa.get_interspike_intervals()