diff --git a/datashuttle/tui/screens/modal_dialogs.py b/datashuttle/tui/screens/modal_dialogs.py index d2fc3dac7..82d7da21d 100644 --- a/datashuttle/tui/screens/modal_dialogs.py +++ b/datashuttle/tui/screens/modal_dialogs.py @@ -13,11 +13,20 @@ from datashuttle.tui.app import TuiApp from datashuttle.utils.custom_types import InterfaceOutput, Prefix +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 ( @@ -201,6 +210,12 @@ def compose(self) -> ComposeResult: yield Container( Static(label_message, id="select_directory_tree_screen_label"), + Select( + [(drive, drive) for drive in self.get_drives()], + value=self.get_selected_drive(), + allow_blank=False, + id="select_directory_tree_drive_select", + ), CustomDirectoryTree( self.mainwindow, self.path_, @@ -210,6 +225,48 @@ def compose(self) -> ComposeResult: id="select_directory_tree_container", ) + @staticmethod + def get_drives(): + """ + Get drives available on the machine to switch between. + For Windows, use `psutil` to get the list of drives. + Otherwise, assume root is "/" and take all folders from that level. + """ + operating_system = platform.system() + + assert operating_system in [ + "Windows", + "Darwin", + "Linux", + ], f"Unexpected operating system: {operating_system} encountered." + + if platform.system() == "Windows": + return [disk.device for disk in psutil.disk_partitions(all=True)] + + else: + return ["/"] + [ + f"/{dir.name}" for dir in Path("/").iterdir() if dir.is_dir() + ] + + def get_selected_drive(self): + """ + Get the default drive which the select starts on. For windows, + use the .drive attribute but for macOS and Linux this is blank. + On these Os use the first folder (e.g. /Users) as the default drive. + """ + if platform.system() == "Windows": + selected_drive = f"{self.path_.drive}\\" + else: + selected_drive = f"/{self.path_.parts[1]}" + return selected_drive + + def on_select_changed(self, event: Select.Changed) -> None: + """Updates the directory tree when the drive is changed.""" + self.path_ = Path(event.value) + self.query_one("#select_directory_tree_directory_tree").path = ( + self.path_ + ) + @require_double_click def on_directory_tree_directory_selected( self, event: DirectoryTree.DirectorySelected diff --git a/pyproject.toml b/pyproject.toml index 9ca0f69e6..528693b54 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -22,7 +22,8 @@ dependencies = [ "textual==3.4.0", "show-in-file-manager", "gitpython", - "typeguard" + "typeguard", + "psutil" ] classifiers = [ diff --git a/tests/tests_tui/test_tui_selectdirectorytree.py b/tests/tests_tui/test_tui_selectdirectorytree.py new file mode 100644 index 000000000..3d1a9d362 --- /dev/null +++ b/tests/tests_tui/test_tui_selectdirectorytree.py @@ -0,0 +1,58 @@ +import pytest +from tui_base import TuiBase + +from datashuttle.tui.app import TuiApp +from datashuttle.tui.screens.modal_dialogs import ( + SelectDirectoryTreeScreen, +) + + +class TestSelectTree(TuiBase): + @pytest.mark.asyncio + async def test_select_directory_tree(self, monkeypatch): + """ + Test that changing the drive in SelectDirectoryTreeScreen + updates the DirectoryTree path as expected. + """ + + # Set the Select drives to be these test cases + monkeypatch.setattr( + SelectDirectoryTreeScreen, + "get_selected_drive", + staticmethod(lambda: "Drive1"), + ) + + monkeypatch.setattr( + SelectDirectoryTreeScreen, + "get_drives", + staticmethod(lambda: ["Drive1", "Drive2"]), + ) + + app = TuiApp() + async with app.run_test() as pilot: + + # Open the select directory tree screen + await self.scroll_to_click_pause( + pilot, "#mainwindow_new_project_button" + ) + + await self.scroll_to_click_pause( + pilot, "#configs_local_path_select_button" + ) + assert isinstance(pilot.app.screen, SelectDirectoryTreeScreen) + + # Switch the select, and ensure the directory tree is updated as expected + tree = pilot.app.screen.query_one( + "#select_directory_tree_directory_tree" + ) + select = pilot.app.screen.query_one( + "#select_directory_tree_drive_select" + ) + + select.value = "Drive1" + await pilot.pause() + assert str(tree.path) == "Drive1" + + select.value = "Drive2" + await pilot.pause() + assert str(tree.path) == "Drive2"