[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

@@ -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

pyproject.toml

@@ -22,7 +22,8 @@ dependencies = [
     "textual==3.4.0",
     "show-in-file-manager",
     "gitpython",
-    "typeguard"
+    "typeguard",
+    "psutil"
 ]
 
 classifiers = [

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"