BerryV2/src/widgets/directory_tree_custom.py

from textual.widgets import DirectoryTree, Rule
from textual.widgets.tree import TreeNode
from textual.events import MouseDown

from prompt import Prompt
from context_menu import ContextMenu, NoSelectStatic

import os, shutil


class CustomDirectoryTree(DirectoryTree):
    
    ICON_NODE = " "
    ICON_NODE_EXPANDED = " "
    ICON_FILE = " "
    
    def __init__(self, path, *, name = None, id = None, classes = None, disabled = False):
        super().__init__(path, name=name, id=id, classes=classes, disabled=disabled)
        self.right_clicked_node: TreeNode | None = None

    

    def context_menu_chosen(self, file_name: str, result: str):
        if result == "Delete":
            def delete_confirm(will_delete: bool | None):
                if will_delete == True:
                    if os.path.isfile(self.right_clicked_node.data.path):
                        os.remove(self.right_clicked_node.data.path)
                    else:
                        shutil.rmtree(self.right_clicked_node.data.path)
                    self.reload()

                    self.notify(f"Deleted \"{self.right_clicked_node.data.path}\".")
                self.right_clicked_node = None

            self.app.push_screen(Prompt(f"Are you sure you want to delete \"{self.right_clicked_node.label}\"?", "confirm", "Confirm deletion"), delete_confirm)
        elif result == "Rename":
            def rename_confirm(new_name: str | None):
                if new_name == None: return
                if new_name.strip() == "":
                    self.notify("Filename can't be empty.", title="Failed to rename", severity="error")
                    return

                os.rename(self.right_clicked_node.data.path, os.path.join(os.path.dirname(self.right_clicked_node.data.path), new_name))
                self.reload()

                self.notify("Renamed successfully.")

                self.right_clicked_node = None

            self.app.push_screen(Prompt(f"Enter the new name for \"{self.right_clicked_node.label}\".", "string", "Rename"), rename_confirm)
        elif result == "New folder":
            def new_folder(folder_name: str | None):
                if folder_name == None: return
                if folder_name.strip() == "":
                    self.notify("Folder name can't be empty.", title="Failed to create folder", severity="error")
                    return                

                new_folder_path = os.path.join(self.right_clicked_node.data.path, folder_name)
                try:
                    os.mkdir(new_folder_path)
                    self.reload()
                except Exception as e:
                    self.notify(str(e), title="Failed to create folder", severity="error")
                    return

            self.app.push_screen(Prompt("Enter the name of the new folder.", "string", "Create New Folder"), new_folder) 
        elif result == "New file":
            def new_file(file_name: str | None):
                if file_name == None: return
                if file_name.strip() == "":
                    self.notify("File name can't be empty.", title="Failed to create file", severity="error")
                    return
                
                new_file_path = os.path.join(self.right_clicked_node.data.path, file_name)

                if os.path.isfile(new_file_path):
                    self.notify("That file already exists.", title="Failed to create file", severity="error")
                    return

                try:
                    with open(new_file_path, "w") as f:
                        pass
                    self.reload()
                except Exception as e:
                    self.notify(str(e), title="Failed to create file", severity="error")
                    return

            self.app.push_screen(Prompt("Enter the name of the new file", "string", "Create New File"), new_file)
        elif result == "Open with...":
            file_extension = file_name.rsplit(".")[1]
            default_editors = eval(self.app.config_handler.get("editor", "default_editors"))
            default_editor = default_editors.get(file_extension, "Text Editor")
            
            def open_with(chosen_editor):
                self.notify(str(chosen_editor))
            
            self.app.push_screen(Prompt(
                "Editor for file:",
                "list",
                "Choose editor for file",
                values=[
                    default_editor + " (Default)",
                    "Text Editor"
                ],
                allow_blank=False
            ), open_with)

    def on_mouse_down(self, event: MouseDown):
        if event.button != 3 or not "line" in event.style.meta:
            return
        selected_node = self.get_node_at_line(event.style.meta["line"])
        self.right_clicked_node = selected_node

        spacer = NoSelectStatic(f'[d]{"-" * 17}[/]')

        options = None
        if self._safe_is_dir(self.right_clicked_node.data.path):
            options = ["New folder", "New file", spacer, "Delete", "Rename"]
        else:
            options = ["Delete", "Rename", spacer, "Open with..."]

        file_name = str(self.right_clicked_node.label) if len(self.right_clicked_node.label) <= 17 else self.right_clicked_node.label[:14] + "..."
        
        self.app.push_screen(ContextMenu(
            [NoSelectStatic(f"[b]{file_name}[/]"), NoSelectStatic(f'[d]{"-" * 17}[/]')] + options,
            event.screen_offset
        ), lambda result: self.context_menu_chosen(file_name, result))