SpaceSniffer/SpaceSniffer.py

import os
import threading
import tkinter as tk
from tkinter import ttk, filedialog, messagebox
from queue import Queue
from send2trash import send2trash
import subprocess


class FolderScanner(threading.Thread):
    def __init__(self, start_path, queue, progress_queue):
        threading.Thread.__init__(self)
        self.start_path = start_path
        self.queue = queue
        self.progress_queue = progress_queue
        self.total_items = 0

    def run(self):
        self.total_items = sum(
            [len(files) + len(dirs) for r, dirs, files in os.walk(self.start_path)]
        )
        self.scan_folder(self.start_path)
        self.queue.put(None)  # Indicate that scanning is done

    def scan_folder(self, folder):
        scanned_items = 0
        for dirpath, dirnames, filenames in os.walk(folder):
            for dirname in dirnames:
                folder_path = os.path.join(dirpath, dirname)
                size = self.get_size(folder_path)
                self.queue.put((folder_path, size, dirpath, "folder"))
                scanned_items += 1
                self.update_progress(scanned_items, folder_path)

            for filename in filenames:
                file_path = os.path.join(dirpath, filename)
                size = os.path.getsize(file_path)
                self.queue.put((file_path, size, dirpath, "file"))
                scanned_items += 1
                self.update_progress(scanned_items, file_path)

    def get_size(self, path):
        total_size = 0
        for dirpath, dirnames, filenames in os.walk(path):
            for f in filenames:
                fp = os.path.join(dirpath, f)
                total_size += os.path.getsize(fp)
        return total_size

    def update_progress(self, scanned_items, current_path):
        progress = (scanned_items / self.total_items) * 100
        self.progress_queue.put(
            (progress, scanned_items, self.total_items, current_path)
        )


class App(tk.Tk):
    def __init__(self, start_path):
        super().__init__()
        self.withdraw()
        self.title("目录扫描器 V1.0")
        self.start_path = start_path
        self.queue = Queue()
        self.progress_queue = Queue()
        self.progress_window = ProgressWindow(self)

        self.scanner = FolderScanner(start_path, self.queue, self.progress_queue)
        self.scanner.start()

        self.style = ttk.Style(self)
        self.style.configure("Treeview", font=("Helvetica", 10), rowheight=25)
        self.style.configure("Treeview.Heading", font=("Helvetica", 12, "bold"))

        self.tree = ttk.Treeview(self, style="Treeview")
        self.tree.heading("#0", text="目录/文件名（目录/文件大小）", anchor="w")
        self.tree.tag_configure("folder", foreground="orange")
        self.tree.tag_configure("file", foreground="black")
        self.tree.pack(fill=tk.BOTH, expand=True)

        self.folder_first_var = tk.BooleanVar()
        self.folder_first_checkbutton = tk.Checkbutton(
            self, text="目录优先显示", variable=self.folder_first_var
        )
        self.folder_first_checkbutton.pack(side=tk.LEFT, fill=tk.X)

        self.folder_first_checkbutton.config(command=self.update_folder_first_var)

        self.sort_by_size_button = tk.Button(
            self, text="根据目录/文件大小排序", command=self.sort_by_size
        )
        self.sort_by_size_button.pack(side=tk.LEFT, fill=tk.X)

        self.sort_by_name_button = tk.Button(
            self, text="根据目录/文件名排序", command=self.sort_by_name
        )
        self.sort_by_name_button.pack(side=tk.LEFT, fill=tk.X)

        self.refresh_button = tk.Button(
            self, text="刷新列表", command=self.refresh_tree
        )
        self.refresh_button.pack(side=tk.LEFT, fill=tk.X)

        self.tree.bind("<Button-3>", self.show_context_menu)

        self.size_sort_order = True
        self.name_sort_order = True
        self.is_sorting = False

        self.populate_root()
        self.update_tree()
        self.update_progress()

        self.protocol("WM_DELETE_WINDOW", self.on_close)
        self.set_window_size()

    def populate_root(self):
        size = self.scanner.get_size(self.start_path) / 1024 / 1024
        self.tree.insert(
            "",
            "end",
            iid=self.start_path,
            text=f"{os.path.basename(self.start_path)} ({size:.2f} MB)",
            tags=("folder",),
            open=True,
        )

    def update_tree(self):
        if not self.is_sorting:
            while not self.queue.empty():
                item = self.queue.get()
                if item is None:  # Scanning is done
                    self.progress_window.destroy()
                    self.deiconify()
                    return

                path, size, parent, tag = item
                size_mb = size / 1024 / 1024
                parent_iid = self.get_iid(parent)
                if parent == self.start_path:
                    parent_iid = self.start_path
                if not self.tree.exists(self.get_iid(path)):
                    self.tree.insert(
                        parent_iid,
                        "end",
                        iid=self.get_iid(path),
                        text=f"{os.path.basename(path)} ({size_mb:.2f} MB)",
                        tags=(tag,),
                        open=False,
                    )
            self.after(100, self.update_tree)

    def update_progress(self):
        if not self.progress_window.winfo_exists():  # Check if the window still exists
            return
        while not self.progress_queue.empty():
            progress, scanned_items, total_items, current_path = (
                self.progress_queue.get()
            )
            self.progress_window.progress_var.set(progress)
            self.progress_window.progress_bar["value"] = progress
            self.progress_window.progress_label.config(
                text=f"扫描进度: {scanned_items} / {total_items}"
            )
            self.progress_window.path_label.config(text=f"正在扫描: {current_path}")

        self.after(100, self.update_progress)

    def get_iid(self, path):
        return path

    def sort_by_size(self):
        self.sort_items(
            lambda x: (
                0 if x[1] == "folder" and self.folder_first_var.get() else 1,
                (
                    -float(x[0].split()[-2].replace("(", ""))
                    if self.size_sort_order
                    else float(x[0].split()[-2].replace("(", ""))
                ),
            ),
            False,
        )
        self.size_sort_order = not self.size_sort_order

    def sort_by_name(self):
        self.sort_items(
            lambda x: (
                0 if x[1] == "folder" and self.folder_first_var.get() else 1,
                x[0].lower(),
            ),
            self.name_sort_order,
        )
        self.name_sort_order = not self.name_sort_order

    def sort_items(self, key, reverse):
        if self.folder_first_var.get():
            selected = self.tree.selection()
            if not selected:
                messagebox.showinfo("提示", "请选择你要排序的目录")
                return

            self.scanner.join()
            self.is_sorting = True

            for parent in selected:
                children = self.tree.get_children(parent)
                folder_items = []
                file_items = []

                for child in children:
                    item = (
                        self.tree.item(child)["text"],
                        self.tree.item(child)["tags"][0],
                        child,
                    )
                    if item[1] == "folder":
                        folder_items.append(item)
                    else:
                        file_items.append(item)

                sorted_items = sorted(folder_items, key=key, reverse=reverse) + sorted(
                    file_items, key=key, reverse=reverse
                )

                for item in sorted_items:
                    self.tree.detach(item[2])
                for index, item in enumerate(sorted_items):
                    self.tree.move(item[2], parent, index)
        else:
            selected = self.tree.selection()
            if not selected:
                messagebox.showinfo("提示", "请选择你要排序的目录")
                return

            self.scanner.join()
            self.is_sorting = True

            for parent in selected:
                children = self.tree.get_children(parent)
                items = [
                    (
                        self.tree.item(child)["text"],
                        self.tree.item(child)["tags"][0],
                        child,
                    )
                    for child in children
                ]
                sorted_items = sorted(items, key=key, reverse=reverse)

                for item in sorted_items:
                    self.tree.detach(item[2])
                for index, item in enumerate(sorted_items):
                    self.tree.move(item[2], parent, index)

    def show_context_menu(self, event):
        item = self.tree.identify_row(event.y)
        if item:
            self.tree.selection_set(item)
            menu = tk.Menu(self, tearoff=0)
            menu.add_command(label="打开文件", command=lambda: self.open_item(item))
            menu.add_command(
                label="打开所在目录", command=lambda: self.open_directory(item)
            )
            menu.add_command(
                label="删除文件/目录(放入回收站)",
                command=lambda: self.delete_item(item),
            )
            menu.post(event.x_root, event.y_root)

    def open_item(self, item):
        path = self.tree.item(item, "text").split(" ")[0]
        full_path = os.path.join(self.start_path, path)
        try:
            if os.path.isdir(full_path):
                subprocess.Popen(f'explorer "{os.path.realpath(full_path)}"')
            else:
                os.startfile(full_path)
        except Exception as e:
            messagebox.showerror("错误", f"无法打开 {full_path}\n{e}")

    def open_directory(self, item):
        path = self.tree.item(item, "text").split(" ")[0]
        full_path = os.path.join(self.start_path, path)
        if not os.path.isdir(full_path):
            full_path = os.path.dirname(full_path)
        try:
            subprocess.Popen(f'explorer "{os.path.realpath(full_path)}"')
        except Exception as e:
            messagebox.showerror("错误", f"无法打开 {full_path}\n{e}")

    def delete_item(self, item):
        path = self.tree.item(item, "text").split(" ")[0]
        full_path = os.path.normpath(os.path.join(self.start_path, path))
        confirm = messagebox.askyesno("提示", f"你确定要删除 {full_path} 吗?")
        if confirm:
            try:
                send2trash(full_path)
                self.tree.delete(item)
            except Exception as e:
                messagebox.showerror("错误", f"无法删除 {full_path}\n{e}")
            self.refresh_tree()

    def refresh_tree(self):
        for item in self.tree.get_children():
            self.tree.delete(item)
        self.populate_root()
        self.scanner = FolderScanner(self.start_path, self.queue, self.progress_queue)
        self.progress_window = ProgressWindow(self)  # Create a new progress window
        self.withdraw()  # Hide the main window
        self.scanner.start()
        self.update_tree()  # Start updating the tree again
        self.update_progress()  # Start updating the progress again

    def on_close(self):
        if self.scanner.is_alive():
            self.scanner.join()
        self.destroy()
        root.destroy()

    def set_window_size(self):
        screen_width = self.winfo_screenwidth()
        screen_height = self.winfo_screenheight()
        window_width = int(screen_width * 0.6)
        window_height = int(screen_height * 0.6)
        self.geometry(
            f"{window_width}x{window_height}+{screen_width//2 - window_width//2}+{screen_height//2 - window_height//2}"
        )
        self.update_font_size()

    def update_font_size(self):
        current_font = self.style.lookup("Treeview", "font")
        font_size = int(current_font.split()[1])
        screen_width = self.winfo_screenwidth()
        new_font_size = int(screen_width / 100)
        if new_font_size != font_size:
            new_font = (current_font.split()[0], new_font_size)
            self.style.configure("Treeview", font=new_font)
            self.style.configure(
                "Treeview.Heading",
                font=(current_font.split()[0], new_font_size + 2, "bold"),
            )

    def update_folder_first_var(self):
        self.folder_first_var.set(not self.folder_first_var.get())


class ProgressWindow(tk.Toplevel):
    def __init__(self, parent):
        super().__init__(parent)
        self.title("扫描进度")
        self.geometry("500x200")
        self.label = tk.Label(
            self, text="扫描中，请稍等", fg="deepskyblue", font=("Helvetica", 16)
        )
        self.label.pack(pady=20)
        self.progress_var = tk.DoubleVar()
        self.progress_bar = ttk.Progressbar(
            self, variable=self.progress_var, maximum=100
        )
        self.progress_bar.pack(fill=tk.X, padx=20, pady=10)
        self.progress_label = tk.Label(self, text="扫描进度: 0 / 0")
        self.progress_label.pack(pady=5)
        self.path_label = tk.Label(self, text="正在扫描: ")
        self.path_label.pack(pady=5)
        self.protocol("WM_DELETE_WINDOW", self.on_close)

    def on_close(self):
        pass  # Prevent the user from closing the progress window


if __name__ == "__main__":
    root = tk.Tk()
    root.withdraw()
    start_path = filedialog.askdirectory()
    if start_path:
        app = App(start_path)
        app.mainloop()
    else:
        messagebox.showinfo("提示", "没有选中任何目录，程序即将退出")
        root.destroy()