From 60a1e080890cb7a004a19a5c7f0cc540a2929591 Mon Sep 17 00:00:00 2001
From: chrishalcrow <57948917+chrishalcrow@users.noreply.github.com>
Date: Wed, 6 Mar 2024 16:50:44 +0000
Subject: [PATCH] Make compute_correlograms compatible with
 MockWaveformExtractor

---
 src/spikeinterface/postprocessing/correlograms.py | 6 ++++++
 1 file changed, 6 insertions(+)

diff --git a/src/spikeinterface/postprocessing/correlograms.py b/src/spikeinterface/postprocessing/correlograms.py
index f826cf9e8d..34da1ffcd4 100644
--- a/src/spikeinterface/postprocessing/correlograms.py
+++ b/src/spikeinterface/postprocessing/correlograms.py
@@ -4,6 +4,8 @@
 import numpy as np
 from spikeinterface.core.sortinganalyzer import register_result_extension, AnalyzerExtension, SortingAnalyzer
 
+from spikeinterface.core.waveforms_extractor_backwards_compatibility import MockWaveformExtractor
+
 try:
     import numba
 
@@ -87,6 +89,10 @@ def compute_correlograms(
     bin_ms: float = 1.0,
     method: str = "auto",
 ):
+    
+    if isinstance(sorting_analyzer_or_sorting, MockWaveformExtractor):
+        sorting_analyzer_or_sorting = sorting_analyzer_or_sorting.sorting
+
     if isinstance(sorting_analyzer_or_sorting, SortingAnalyzer):
         return compute_correlograms_sorting_analyzer(
             sorting_analyzer_or_sorting, window_ms=window_ms, bin_ms=bin_ms, method=method