Berry/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):
	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, result):
		if result == "Open":
			self.select_node(self.right_clicked_node)
			self.right_clicked_node = None
		elif result == "Delete":
			def delete_confirm(will_delete: bool | None):
				if will_delete == True:
					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

				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)

	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

		options = None
		if self._safe_is_dir(self.right_clicked_node.data.path):
			options = ["New Folder", "New File", NoSelectStatic(f'[d]{"-" * 17}[/]'), "Delete", "Rename", "Open"]
		else:
			options = ["Delete", "Rename", "Open"]
		
		self.app.push_screen(ContextMenu(
			[NoSelectStatic(f"[b]{self.right_clicked_node.label}[/]"), NoSelectStatic(f'[d]{"-" * 17}[/]')] + options,
			event.screen_offset
		), self.context_menu_chosen)