ui.workspace_tree_view

  1import os
  2import re
  3
  4from PyQt6 import QtWidgets, QtCore
  5from PyQt6.QtCore import QUrl, pyqtSignal
  6from PyQt6.QtWidgets import QTreeView, QMessageBox, QMenu, QFileDialog
  7from PyQt6.QtGui import QFileSystemModel, QDesktopServices, QAction
  8
  9from components.file_role_dialog import FileRoleDialog
 10from logger import get_logger
 11from threads.utils_threads import CopyDeleteThread
 12
 13log = get_logger()
 14
 15
 16class WorkspaceTreeView(QTreeView):
 17    """
 18    Custom QTreeView to manage files and folders within a workspace.
 19
 20    This widget:
 21      - Displays the workspace directory structure using QFileSystemModel.
 22      - Handles context menus for file/folder operations (add, remove, export, open).
 23      - Supports multiple file selection and drag/drop-like functionality via threads.
 24      - Integrates with the main application context for shared signals and callbacks.
 25    """
 26
 27    # Signal emitted when a background thread (Copy/Delete) is created
 28    new_thread = pyqtSignal(object)
 29    """**Signal(object):** Emitted when a new background thread (e.g., for copy or delete operations) is created.  
 30    Parameter:
 31    - `object`: Reference to the created thread instance.
 32    """
 33
 34    def __init__(self, context):
 35        super().__init__()
 36
 37        # Context dictionary (shared among pages/windows)
 38        self.context = context
 39        self.workspace_path = self.context["workspace_path"]
 40
 41        self.selected_files = []
 42
 43        # --- File System Model setup ---
 44        self.tree_model = QFileSystemModel()
 45        self.tree_model.setRootPath(self.workspace_path)
 46
 47        # Apply model to the tree view
 48        self.setModel(self.tree_model)
 49        self.setRootIndex(self.tree_model.index(self.workspace_path))
 50        self.setMinimumSize(QtCore.QSize(200, 0))
 51
 52        # Enable extended selection (multi-selection)
 53        self.setSelectionMode(QtWidgets.QAbstractItemView.SelectionMode.ExtendedSelection)
 54
 55        # Enable custom context menu
 56        self.setContextMenuPolicy(QtCore.Qt.ContextMenuPolicy.CustomContextMenu)
 57
 58        # --- Signal connections ---
 59        self.clicked.connect(self.handle_workspace_click)
 60        self.doubleClicked.connect(self.handle_double_click)
 61        self.customContextMenuRequested.connect(self.open_tree_context_menu)
 62
 63    # ---------------------------------------------------------------
 64    # Event Handlers
 65    # ---------------------------------------------------------------
 66
 67    def handle_workspace_click(self):
 68        """Handle single-click event: update selected files and emit them via context signal."""
 69        selected_indexes = self.selectionModel().selectedRows()
 70        selected_files = [self.tree_model.filePath(idx) for idx in selected_indexes]
 71
 72        self.selected_files = selected_files
 73        self.context["selected_files_signal"].emit(selected_files)
 74
 75    def handle_double_click(self, index):
 76        """
 77        Handle double-click event:
 78        - If the selected item is a NIfTI file (.nii / .nii.gz), open it in the viewer.
 79        - Otherwise, open it with the system's default application.
 80        """
 81        file_path = self.tree_model.filePath(index)
 82        if not os.path.isfile(file_path):
 83            return
 84
 85        if file_path.endswith((".nii", ".nii.gz")):
 86            try:
 87                self.context["open_nifti_viewer"](file_path)
 88            except Exception as e:
 89                QMessageBox.critical(
 90                    self,
 91                    QtCore.QCoreApplication.translate("TreeView", "Error"),
 92                    QtCore.QCoreApplication.translate("TreeView", "Error when opening NIfTI file:\n{0}").format(str(e))
 93                )
 94                log.error(f"Error opening NIfTI file:\n{str(e)}")
 95        else:
 96            self._open_in_explorer(file_path)
 97
 98    def adjust_tree_columns(self):
 99        """Automatically hide or show extra columns based on the widget width."""
100        width = self.width()
101        for i in range(1, self.tree_model.columnCount()):
102            if width > 400:
103                self.showColumn(i)
104                self.setColumnWidth(i, 100)
105            else:
106                self.hideColumn(i)
107
108    # ---------------------------------------------------------------
109    # Context Menu
110    # ---------------------------------------------------------------
111
112    def open_tree_context_menu(self, position):
113        """
114        Open the right-click context menu dynamically, depending on the selected item:
115        - Workspace root
116        - Folder
117        - File
118        - Multiple files
119        """
120        index = self.indexAt(position)
121        file_path = None
122        is_dir = False
123        is_nifty = False
124
125        # Determine selected file/folder
126        if self.selected_files and len(self.selected_files) == 1 and not index.isValid():
127            file_path = self.selected_files[0]
128            is_dir = self.tree_model.isDir(index) if index.isValid() else False
129            is_nifty = file_path.endswith((".nii", ".nii.gz"))
130        elif index.isValid():
131            file_path = self.tree_model.filePath(index)
132            is_dir = self.tree_model.isDir(index)
133
134        menu = QMenu(self)
135        actions = {}
136
137        # Choose menu layout based on selection
138        if not index.isValid() and not self.selected_files:
139            actions.update(self._add_workspace_actions(menu))
140        elif index.isValid() and file_path and len(self.selected_files) <= 1:
141            if is_dir:
142                actions.update(self._add_folder_actions(menu, file_path))
143            else:
144                actions.update(self._add_file_actions(menu, file_path, is_nifty))
145        elif index.isValid() and len(self.selected_files) > 1:
146            actions.update(self._add_multi_file_actions(menu))
147
148        # Execute chosen action
149        chosen_action = menu.exec(self.viewport().mapToGlobal(position))
150        if chosen_action and chosen_action in actions:
151            actions[chosen_action](file_path, is_dir, is_nifty)
152
153    # --- Context Menu Definitions ---
154
155    def _add_workspace_actions(self, menu):
156        """Add actions available at the root workspace level."""
157        open_in_explorer = menu.addAction(QtCore.QCoreApplication.translate("TreeView", "Open workspace in explorer"))
158        menu.addSeparator()
159        export_workspace = menu.addAction(QtCore.QCoreApplication.translate("TreeView", "Export workspace"))
160        clear_workspace = menu.addAction(QtCore.QCoreApplication.translate("TreeView", "Clear workspace"))
161        menu.addSeparator()
162        add_single_file = menu.addAction(QtCore.QCoreApplication.translate("TreeView", "Add single file to workspace"))
163
164        return {
165            open_in_explorer: lambda *_: self._open_in_explorer(self.workspace_path),
166            export_workspace: lambda *_: self.export_files([self.workspace_path], True),
167            clear_workspace: lambda *_: self.context["main_window"].clear_folder(
168                folder_path=self.workspace_path, folder_name="workspace", return_to_import=True
169            ),
170            add_single_file: lambda *_: self.add_file_to_workspace(None, False),
171        }
172
173    def _add_folder_actions(self, menu, file_path):
174        """Add actions for a folder."""
175        open_action = menu.addAction(QtCore.QCoreApplication.translate("TreeView", "Open folder in explorer"))
176        menu.addSeparator()
177        add_action = menu.addAction(QtCore.QCoreApplication.translate("TreeView", "Add single file to folder"))
178        remove_action = menu.addAction(QtCore.QCoreApplication.translate("TreeView", "Remove folder from workspace"))
179        export_action = menu.addAction(QtCore.QCoreApplication.translate("TreeView", "Export folder"))
180
181        return {
182            open_action: lambda *_: self._open_in_explorer(file_path),
183            add_action: lambda *_: self.add_file_to_workspace(file_path, True),
184            remove_action: lambda *_: self.remove_from_workspace([file_path]),
185            export_action: lambda *_: self.export_files([file_path], True),
186        }
187
188    def _add_file_actions(self, menu, file_path, is_nifty):
189        """Add actions for a single file."""
190        open_action = menu.addAction(QtCore.QCoreApplication.translate("TreeView", "Open with system predefined"))
191        actions = {open_action: lambda *_: self._open_in_explorer(file_path)}
192
193        if is_nifty:
194            nifti_action = menu.addAction(QtCore.QCoreApplication.translate("TreeView", "Open NIfTI file"))
195            actions[nifti_action] = lambda *_: self._open_nifti(file_path)
196
197        menu.addSeparator()
198        add_action = menu.addAction(QtCore.QCoreApplication.translate("TreeView", "Add single file"))
199        remove_action = menu.addAction(QtCore.QCoreApplication.translate("TreeView", "Remove file from workspace"))
200        export_action = menu.addAction(QtCore.QCoreApplication.translate("TreeView", "Export file"))
201
202        actions.update({
203            add_action: lambda *_: self.add_file_to_workspace(file_path, False),
204            remove_action: lambda *_: self.remove_from_workspace([file_path]),
205            export_action: lambda *_: self.export_files([file_path], False),
206        })
207        return actions
208
209    def _add_multi_file_actions(self, menu):
210        """Add actions for multiple selected files."""
211        export_action = menu.addAction(QtCore.QCoreApplication.translate("TreeView", "Export files"))
212        remove_action = menu.addAction(QtCore.QCoreApplication.translate("TreeView", "Remove Files from workspace"))
213        return {
214            export_action: lambda *_: self.export_files(self.selected_files, False),
215            remove_action: lambda *_: self.remove_from_workspace(self.selected_files),
216        }
217
218    # ---------------------------------------------------------------
219    # Core Operations
220    # ---------------------------------------------------------------
221
222    def add_file_to_workspace(self, folder_path, is_dir):
223        """
224        Add a single file into the workspace.
225        Opens a file dialog, optionally adds related JSON, and triggers background copy thread.
226        """
227        dialog = QFileDialog(self.context['main_window'], "Select File")
228        dialog.setFileMode(QFileDialog.FileMode.ExistingFile)
229        dialog.setOption(QFileDialog.Option.DontUseNativeDialog, True)
230        dialog.setOption(QFileDialog.Option.ReadOnly, True)
231        dialog.setDirectory(os.path.expanduser("~"))
232
233        if not dialog.exec():
234            return
235
236        file = dialog.selectedFiles()[0]
237        file_name = os.path.basename(file)
238        dir_name = os.path.dirname(file)
239        name, ext = os.path.splitext(file_name)
240        json_file = None
241
242        # Handle potential .nii.gz format
243        if ext == ".gz":
244            name2, ext2 = os.path.splitext(name)
245            if ext2 == ".nii":
246                json_file = self._including_json(name2, dir_name)
247        elif ext == ".nii":
248            json_file = self._including_json(name, dir_name)
249
250        # Determine destination folder
251        if folder_path:
252            if is_dir:
253                folder = os.path.basename(folder_path)
254            else:
255                folder = os.path.dirname(folder_path)
256        else:
257            folder = "workspace"
258
259        # Specialized behavior based on folder type
260        if folder in ("anat", "pet"):
261            self.new_thread.emit(CopyDeleteThread(src=file, dst=folder_path, is_folder=is_dir, copy=True))
262        elif re.match(r"^ses-\d+$", folder):
263            self.new_thread.emit(CopyDeleteThread(src=file, dst=os.path.join(folder_path, "pet"), is_folder=is_dir, copy=True))
264        elif re.match(r"^sub-\d+$", folder):
265            return self.open_role_dialog(files=[file, json_file], folder_path=folder_path, subj=folder)
266        elif folder == "derivatives":
267            return self.open_role_dialog(files=[file, json_file], folder_path=folder_path, main=folder)
268        else:
269            return self.open_role_dialog(files=[file, json_file], folder_path=self.workspace_path)
270
271    def open_role_dialog(self, files, folder_path=None, subj=None, role=None, main=None):
272        """Open a FileRoleDialog to determine where to place files in the workspace hierarchy."""
273        dialog = FileRoleDialog(workspace_path=self.workspace_path, subj=subj, role=role, main=main, parent=self)
274        if dialog.exec():
275            relative_path = dialog.get_relative_path()
276            path = os.path.join(folder_path, relative_path)
277            os.makedirs(path, exist_ok=True)
278            for file in files:
279                if file:
280                    self.new_thread.emit(CopyDeleteThread(src=file, dst=path, is_folder=False, copy=True))
281
282    def _including_json(self, file_name, file_dir):
283        """
284        Ask the user whether to include an associated JSON file when adding a NIfTI file.
285        Returns JSON path if available and confirmed, otherwise empty string.
286        """
287        answer = QMessageBox.information(
288            self,
289            QtCore.QCoreApplication.translate("TreeView", "Adding Json?"),
290            QtCore.QCoreApplication.translate("TreeView", "Do you want to include the JSON file if present?"),
291            QMessageBox.StandardButton.Ok | QMessageBox.StandardButton.Cancel
292        )
293        if answer == QMessageBox.StandardButton.Ok:
294            json_file_path = os.path.join(file_dir, file_name + ".json")
295            if os.path.isfile(json_file_path):
296                return json_file_path
297        return ""
298
299    # ---------------------------------------------------------------
300    # Export / Remove / Open
301    # ---------------------------------------------------------------
302
303
304    def export_files(self, paths, is_dir):
305        """Export selected files or folders to a chosen destination."""
306        if not paths:
307            return
308
309        def _get_file_filter(ext):
310            """Return the QFileDialog filter based on file extension."""
311            if ext == ".nii.gz":
312                return "NIfTI compressed (*.nii.gz);;NIfTI (*.nii);;All files (*.*)"
313            elif ext == ".nii":
314                return "NIfTI (*.nii);;All files (*.*)"
315            elif ext == ".json":
316                return "JSON (*.json);;All files (*.*)"
317            return "All files (*.*)"
318
319        def _get_json_file(path):
320            """Return associated JSON file if it exists."""
321            base, ext = os.path.splitext(path)
322            if ext == ".gz" and base.endswith(".nii"):
323                base2 = base[:-4]  # remove .nii
324                json_file = os.path.join(os.path.dirname(path), base2 + ".json")
325                return json_file if os.path.exists(json_file) else None
326            elif ext == ".nii":
327                json_file = os.path.join(os.path.dirname(path), base + ".json")
328                return json_file if os.path.exists(json_file) else None
329            return None
330
331        if len(paths) == 1:
332            path = paths[0]
333            if is_dir:
334                # Choose destination folder for a single directory
335                dst_path = QFileDialog.getExistingDirectory(self, "Select destination folder")
336                if dst_path:
337                    dst_path = os.path.join(dst_path, os.path.basename(path))
338                    self.new_thread.emit(CopyDeleteThread(src=path, dst=dst_path, is_folder=True, copy=True))
339            else:
340                # Single file export
341                file_filter = _get_file_filter(os.path.splitext(path)[1])
342                dst_path, _ = QFileDialog.getSaveFileName(self, "Save file as...", os.path.basename(path), file_filter)
343                if not dst_path:
344                    return
345
346                self.new_thread.emit(CopyDeleteThread(src=path, dst=dst_path, is_folder=False, copy=True))
347
348                # Copy associated JSON if exists
349                json_file = _get_json_file(path)
350                if json_file:
351                    dst_json = os.path.splitext(dst_path)[0] + ".json"
352                    self.new_thread.emit(CopyDeleteThread(src=json_file, dst=dst_json, is_folder=False, copy=True))
353        else:
354            # Multiple files/folders
355            dst_dir = QFileDialog.getExistingDirectory(self, "Select destination folder")
356            if not dst_dir:
357                return
358
359            for path in paths:
360                if not os.path.exists(path):
361                    continue
362
363                is_folder = os.path.isdir(path)
364                dst_path = os.path.join(dst_dir, os.path.basename(path)) if is_folder else os.path.join(dst_dir,
365                                                                                                        os.path.basename(
366                                                                                                            path))
367                self.new_thread.emit(CopyDeleteThread(src=path, dst=dst_path, is_folder=is_folder, copy=True))
368
369    def remove_from_workspace(self, paths):
370        """Ask confirmation before deleting selected files/folders from workspace."""
371        for path in paths:
372            if not os.path.exists(path):
373                return
374        if len(paths)<2:
375            is_dir = os.path.isdir(paths[0])
376            item_type = QtCore.QCoreApplication.translate("TreeView", "this folder: ") if is_dir else QtCore.QCoreApplication.translate("TreeView", "this file:")
377        else: item_type = QtCore.QCoreApplication.translate("TreeView", "these files")
378
379        item_name = os.path.basename(paths[0]) if len(paths) == 1 else ""
380        message = QtCore.QCoreApplication.translate("TreeView", "Are you sure you want to remove {0} {1}?").format(item_type, item_name)
381
382        reply = QMessageBox.question(
383            self,
384            QtCore.QCoreApplication.translate("TreeView", "Confirm"),
385            message,
386            QMessageBox.StandardButton.Yes | QMessageBox.StandardButton.No
387        )
388
389        if reply == QMessageBox.StandardButton.Yes:
390            for path in paths:
391                is_dir = os.path.isdir(path)
392                self.new_thread.emit(CopyDeleteThread(src=path, is_folder=is_dir,delete=True))
393
394    def _open_in_explorer(self, path):
395        """Open a file or folder using the system's file explorer or default app."""
396        try:
397            url = QUrl.fromLocalFile(path)
398            success = QDesktopServices.openUrl(url)
399
400            if not success:
401                # Extract file extension
402                ext = os.path.splitext(path)[1] or QtCore.QCoreApplication.translate("TreeView", "(unknown)")
403                QMessageBox.warning(
404                    self,
405                    QtCore.QCoreApplication.translate("TreeView", "No default application"),
406                    QtCore.QCoreApplication.translate("TreeView", "No default application is registered for files with extension {0}.").format(ext)
407                )
408                log.debug(f"No default application is registered for files with extension {ext}.")
409        except Exception as e:
410            QMessageBox.critical(
411                self,
412                QtCore.QCoreApplication.translate("TreeView", "Error"),
413                QtCore.QCoreApplication.translate("TreeView", "An unexpected error occurred while opening the file:\n{0}").format(e)
414            )
415            log.error(f"An unexpected error occurred while opening the file:\n{e}")
416
417    def _open_nifti(self, file_path):
418        """Open a NIfTI file in the integrated viewer."""
419        try:
420            self.context["open_nifti_viewer"](file_path)
421        except Exception as e:
422            QMessageBox.critical(
423                self,
424                QtCore.QCoreApplication.translate("TreeView", "Error"),
425                QtCore.QCoreApplication.translate("TreeView", "Error opening file NIfTI:\n{0}").format(str(e)))
426            log.error(f"Error opening file NIfTI:\n{str(e)}")