[ENH] Added dropdown to select drives and to select files from the tree view while selecting location of new project

datashuttle/tui/screens/modal_dialogs.py

@@ -12,11 +12,21 @@
     from datashuttle.tui.app import App
     from datashuttle.utils.custom_types import InterfaceOutput
 
+import os
+import platform
 from pathlib import Path
 
+import psutil
 from textual.containers import Container, Horizontal
 from textual.screen import ModalScreen
-from textual.widgets import Button, Input, Label, LoadingIndicator, Static
+from textual.widgets import (
+    Button,
+    Input,
+    Label,
+    LoadingIndicator,
+    Select,
+    Static,
+)
 
 from datashuttle.tui.custom_widgets import CustomDirectoryTree
 from datashuttle.tui.utils.tui_decorators import require_double_click
@@ -159,10 +169,16 @@ def __init__(self, mainwindow: App, path_: Optional[Path] = None) -> None:
         super(SelectDirectoryTreeScreen, self).__init__()
         self.mainwindow = mainwindow
 
+        self.available_drives = self.get_drives()
+
         if path_ is None:
             path_ = Path().home()
         self.path_ = path_
 
+        if platform.system() == "Windows":
+            self.selected_drive = self.path_.drive + "\\"
+        else:
+            self.selected_drive = "/"
         self.prev_click_time = 0
 
     def compose(self) -> ComposeResult:
@@ -174,6 +190,12 @@ def compose(self) -> ComposeResult:
 
         yield Container(
             Static(label_message, id="select_directory_tree_screen_label"),
+            Select(  # Dropdown for drives
+                [(drive, drive) for drive in self.available_drives],
+                value=self.selected_drive,
+                allow_blank=False,
+                id="select_directory_tree_drive_select",
+            ),
             CustomDirectoryTree(
                 self.mainwindow,
                 self.path_,
@@ -183,6 +205,27 @@ def compose(self) -> ComposeResult:
             id="select_directory_tree_container",
         )
 
+    @staticmethod
+    def get_drives():
+        if platform.system() == "Windows":
+            return [disk.device for disk in psutil.disk_partitions(all=False)]
+
+        elif platform.system() in ["Darwin", "Linux"]:
+            return ["/"] + [
+                f"/{d}" for d in os.listdir("/") if os.path.isdir(f"/{d}")
+            ]
+        return ["/"]
+
+    def on_select_changed(self, event: Select.Changed) -> None:
+        """Updates the directory tree when the drive is changed."""
+        print(f"Drive selected: {event.value}")  # Debug message
+
+        self.selected_drive = event.value
+        self.path_ = Path(self.selected_drive)
+        self.query_one("#select_directory_tree_directory_tree").path = (
+            self.path_
+        )
+
     @require_double_click
     def on_directory_tree_directory_selected(self, node) -> None:
         if node.path.is_file():