From 867729102ee5a76f412f1a8e7c025ceefadb7bff Mon Sep 17 00:00:00 2001
From: JoeZiminski <joseph.j.ziminski@gmail.com>
Date: Thu, 27 Jun 2024 09:37:21 +0100
Subject: [PATCH] Don't support 4.0.4

---
 .github/scripts/check_kilosort4_releases.py      | 7 +++++++
 .github/scripts/test_kilosort4_ci.py             | 3 ++-
 .github/workflows/test_kilosort4.yml             | 2 +-
 src/spikeinterface/sorters/external/kilosort4.py | 5 +++++
 4 files changed, 15 insertions(+), 2 deletions(-)

diff --git a/.github/scripts/check_kilosort4_releases.py b/.github/scripts/check_kilosort4_releases.py
index 3d04d6948a..9572f88330 100644
--- a/.github/scripts/check_kilosort4_releases.py
+++ b/.github/scripts/check_kilosort4_releases.py
@@ -6,14 +6,21 @@
 
 
 def get_pypi_versions(package_name):
+    """
+    Make an API call to pypi to retrieve all
+    available versions of the kilosort package.
+    """
     url = f"https://pypi.org/pypi/{package_name}/json"
     response = requests.get(url)
     response.raise_for_status()
     data = response.json()
+    versions = list(sorted(data["releases"].keys()))
+    versions.pop(versions.index("4.0.4"))
     return list(sorted(data["releases"].keys()))
 
 
 if __name__ == "__main__":
+    # Get all KS4 versions from pipi and write to file.
     package_name = "kilosort"
     versions = get_pypi_versions(package_name)
     with open(Path(os.path.realpath(__file__)).parent / "kilosort4-latest-version.json", "w") as f:
diff --git a/.github/scripts/test_kilosort4_ci.py b/.github/scripts/test_kilosort4_ci.py
index 3e74fa708e..c894ed71ff 100644
--- a/.github/scripts/test_kilosort4_ci.py
+++ b/.github/scripts/test_kilosort4_ci.py
@@ -342,7 +342,7 @@ def _check_arguments(self, object_, expected_arguments):
 
     # Full Test ####
     @pytest.mark.parametrize("parameter", PARAMS_TO_TEST)
-    def test_kilosort4(self, recording_and_paths, default_results, tmp_path, parameter):
+    def test_kilosort4_main(self, recording_and_paths, default_results, tmp_path, parameter):
         """
         Given a recording, paths to raw data, and a parameter to change,
         run KS4 natively and within the SpikeInterface wrapper with the
@@ -398,6 +398,7 @@ def test_kilosort4(self, recording_and_paths, default_results, tmp_path, paramet
         if parse(version("kilosort")) > parse("4.0.4"):
             self._check_test_parameters_are_changing_the_output(results, default_results, param_key)
 
+    @pytest.mark.skipif(parse(version("kilosort")) == parse("4.0.9"), reason="nblock=0 fails on KS4=4.0.9")
     def test_kilosort4_no_correction(self, recording_and_paths, tmp_path):
         """
         Test the SpikeInterface wrappers `do_correction` argument. We set
diff --git a/.github/workflows/test_kilosort4.yml b/.github/workflows/test_kilosort4.yml
index 3ad61c0d2e..03db2b6170 100644
--- a/.github/workflows/test_kilosort4.yml
+++ b/.github/workflows/test_kilosort4.yml
@@ -44,7 +44,7 @@ jobs:
     strategy:
       fail-fast: false
       matrix:
-        python-version: ["3.10"]  # TODO: just checking python version is not cause of failing test.
+        python-version: ["3.12"]  # TODO: just checking python version is not cause of failing test.
         os: [ubuntu-latest]
         ks_version: ${{ fromJson(needs.versions.outputs.matrix) }}
     steps:
diff --git a/src/spikeinterface/sorters/external/kilosort4.py b/src/spikeinterface/sorters/external/kilosort4.py
index 811a6e8452..55e694a02f 100644
--- a/src/spikeinterface/sorters/external/kilosort4.py
+++ b/src/spikeinterface/sorters/external/kilosort4.py
@@ -163,6 +163,11 @@ def _run_from_folder(cls, sorter_output_folder, params, verbose):
 
             logging.basicConfig(level=logging.INFO)
 
+        if cls.get_sorter_version() == version.parse("4.0.4"):
+            raise RuntimeError(
+                "Kilosort version 4.0.4 is not supported" "in SpikeInterface. Please change Kilosort version."
+            )
+
         sorter_output_folder = sorter_output_folder.absolute()
 
         probe_filename = sorter_output_folder / "probe.prb"