diff --git a/datashuttle/tui/screens/modal_dialogs.py b/datashuttle/tui/screens/modal_dialogs.py index 97706db0..3d24d04e 100644 --- a/datashuttle/tui/screens/modal_dialogs.py +++ b/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():