Manual close of all memmap BEFORE cleanup

trx/trx_file_memmap.py

@@ -19,9 +19,10 @@
 import numpy as np
 
 from trx.io import get_trx_tmpdir
-from trx.utils import (get_reference_info_wrapper,
+from trx.utils import (append_generator_to_dict,
+                       close_or_delete_mmap,
                        convert_data_dict_to_tractogram,
-                       append_generator_to_dict)
+                       get_reference_info_wrapper)
 
 try:
     import dipy
@@ -1736,6 +1737,22 @@ def to_sft(self, resize=False):
     def close(self) -> None:
         """Cleanup on-disk temporary folder and initialize an empty TrxFile"""
         if self._uncompressed_folder_handle is not None:
+            close_or_delete_mmap(self.streamlines)
+
+            # # Close or delete attributes in dictionaries
+            for key in self.data_per_vertex:
+                close_or_delete_mmap(self.data_per_vertex[key])
+
+            for key in self.data_per_streamline:
+                close_or_delete_mmap(self.data_per_streamline[key])
+
+            for key in self.groups:
+                close_or_delete_mmap(self.groups[key])
+
+            for key in self.data_per_group:
+                for dpg in self.data_per_group[key]:
+                    close_or_delete_mmap(self.data_per_group[key][dpg])
+
             try:
                 self._uncompressed_folder_handle.cleanup()
             except PermissionError:

trx/utils.py

@@ -17,6 +17,28 @@
     dipy_available = False
 
 
+def close_or_delete_mmap(obj):
+    """
+    Close the memory-mapped file if it exists, otherwise set the object to None.
+
+    Parameters:
+    -----------
+    obj : object
+        The object that potentially has a memory-mapped file to be closed.
+
+    """
+    if hasattr(obj, '_mmap') and obj._mmap is not None:
+        obj._mmap.close()
+    elif isinstance(obj, ArraySequence):
+        close_or_delete_mmap(obj._data)
+        close_or_delete_mmap(obj._offsets)
+        close_or_delete_mmap(obj._lengths)
+    elif isinstance(obj, np.memmap):
+        del obj
+    else:
+        logging.warning('Object to be close or deleted must be np.memmap')
+
+
 def split_name_with_gz(filename):
     """
     Returns the clean basename and extension of a file.
@@ -115,8 +137,7 @@ def get_reference_info_wrapper(reference):
         voxel_sizes = header['voxel_sizes']
         voxel_order = header['voxel_order']
     elif is_sft:
-        affine, dimensions, voxel_sizes, voxel_order =\
-            reference.space_attributes
+        affine, dimensions, voxel_sizes, voxel_order = reference.space_attributes
     elif is_trx:
         affine = header['VOXEL_TO_RASMM']
         dimensions = header['DIMENSIONS']
@@ -152,10 +173,10 @@ def is_header_compatible(reference_1, reference_2):
         Does all the spatial attribute match
     """
 
-    affine_1, dimensions_1, voxel_sizes_1, voxel_order_1 = \
-        get_reference_info_wrapper(reference_1)
-    affine_2, dimensions_2, voxel_sizes_2, voxel_order_2 = \
-        get_reference_info_wrapper(reference_2)
+    affine_1, dimensions_1, voxel_sizes_1, voxel_order_1 = get_reference_info_wrapper(
+        reference_1)
+    affine_2, dimensions_2, voxel_sizes_2, voxel_order_2 = get_reference_info_wrapper(
+        reference_2)
 
     identical_header = True
     if not np.allclose(affine_1, affine_2, rtol=1e-03, atol=1e-03):