Repository: SpikeInterface/spikeinterface
Pull request #1452 (2023-03-23 18:09:07)

Add memory to recording `__repr__`
Quick PR as discussed today.

![image](https://user-images.githubusercontent.com/6429509/227305498-e00d2f6b-d689-4f7a-91ab-ed2feaad871e.png)



Patch:
diff --git a/spikeinterface/core/baserecording.py b/spikeinterface/core/baserecording.py
index f1f724f60..aa00f523d 100644
--- a/spikeinterface/core/baserecording.py
+++ b/spikeinterface/core/baserecording.py
@@ -10,6 +10,7 @@ from .base import BaseSegment
 from .baserecordingsnippets import BaseRecordingSnippets
 from .core_tools import write_binary_recording, write_memory_recording, write_traces_to_zarr, check_json
 from .job_tools import split_job_kwargs, fix_job_kwargs
+from .core_tools import convert_bytes_to_str
 
 from warnings import warn
 
@@ -42,12 +43,17 @@ class BaseRecording(BaseRecordingSnippets):
         nchan = self.get_num_channels()
         sf_khz = self.get_sampling_frequency() / 1000.
         duration = self.get_total_duration()
-        txt = f'{clsname}: {nchan} channels - {nseg} segments - {sf_khz:0.1f}kHz - {duration:0.3f}s'
+        memory_size = self.get_memory_size()
+        txt = f"{clsname}: {nchan} channels - {nseg} segments - {sf_khz:0.1f}kHz - {duration:0.3f}s - {memory_size}"
         if 'file_paths' in self._kwargs:
             txt += '\n  file_paths: {}'.format(self._kwargs['file_paths'])
         if 'file_path' in self._kwargs:
             txt += '\n  file_path: {}'.format(self._kwargs['file_path'])
         return txt
+    
+    def get_memory_size(self):
+        bytes = self.get_total_samples() * self.get_num_channels() * self.get_dtype().itemsize 
+        return convert_bytes_to_str(bytes)
 
     def get_num_segments(self):
         """Returns the number of segments.
diff --git a/spikeinterface/core/core_tools.py b/spikeinterface/core/core_tools.py
index 9288a4a51..cb9b43e94 100644
--- a/spikeinterface/core/core_tools.py
+++ b/spikeinterface/core/core_tools.py
@@ -811,3 +811,21 @@ def recursive_key_finder(d, key):
         else:
             if k == key:
                 yield v
+
+
+def convert_bytes_to_str(byte_value:int ) -> str:
+    """
+    Converts a number of bytes to a value in either KiB, MiB, GiB, or TiB.
+
+    Args:
+        byte_value (int): The number of bytes to convert.
+
+    Returns:
+        str: The converted value with the appropriate unit (KiB, MiB, GiB, or TiB).
+    """
+    suffixes = ['B', 'KiB', 'MiB', 'GiB', 'TiB']
+    i = 0
+    while byte_value >= 1024 and i < len(suffixes) - 1:
+        byte_value /= 1024
+        i += 1
+    return f"{byte_value:.2f} {suffixes[i]}"
\ No newline at end of file
diff --git a/spikeinterface/core/generate.py b/spikeinterface/core/generate.py
index e1b473a20..078b1b924 100644
--- a/spikeinterface/core/generate.py
+++ b/spikeinterface/core/generate.py
@@ -516,7 +516,7 @@ def generate_lazy_recording(full_traces_size_GiB: float, seed=None) -> Generator
     num_channels = 1024    
     
     GiB_to_bytes = 1024** 3
-    full_traces_size_bytes = full_traces_size_GiB * GiB_to_bytes 
+    full_traces_size_bytes = int(full_traces_size_GiB * GiB_to_bytes) 
     num_samples = int(full_traces_size_bytes / (num_channels * dtype.itemsize))
     durations = [num_samples / sampling_frequency]
 
@@ -525,6 +525,7 @@ def generate_lazy_recording(full_traces_size_GiB: float, seed=None) -> Generator
 
     return recording
 
+
 if __name__ == '__main__':
     print(generate_recording())
     print(generate_sorting())
diff --git a/spikeinterface/core/job_tools.py b/spikeinterface/core/job_tools.py
index a7edb5746..38f7722e9 100644
--- a/spikeinterface/core/job_tools.py
+++ b/spikeinterface/core/job_tools.py
@@ -397,3 +397,5 @@ def function_wrapper(args):
     else:
         with threadpool_limits(limits=max_threads_per_process):
             return _func(segment_index, start_frame, end_frame, _worker_ctx)
+
+


Test patch:
diff --git a/spikeinterface/core/tests/test_generate.py b/spikeinterface/core/tests/test_generate.py
index e80fd727f..cdbbe638d 100644
--- a/spikeinterface/core/tests/test_generate.py
+++ b/spikeinterface/core/tests/test_generate.py
@@ -98,3 +98,9 @@ def test_generate_lazy_recording():
     print(f"Difference between the last two {(memory_after_traces_MiB - traces_size_MiB)} MiB")
 
     (memory_after_instanciation_MiB + traces_size_MiB) == pytest.approx(memory_after_traces_MiB, rel=relative_tolerance)
+
+
+def test_generate_lazy_recording_under_giga():
+    
+    recording = generate_lazy_recording(full_traces_size_GiB=0.5)
+    assert recording.get_memory_size() == "512.00 MiB"