IamHere/folder_tree_3d.py

import tkinter as tk
from tkinter import filedialog, ttk, messagebox
import os
import math
from tkinter import Canvas

class FolderTree3D:
    def __init__(self, root):
        self.root = root
        self.root.title("文件夹3D树状图")
        self.root.geometry("1200x800")
        
        # 创建主框架
        main_frame = ttk.Frame(root)
        main_frame.pack(fill=tk.BOTH, expand=True, padx=10, pady=10)
        
        # 输入区域
        input_frame = ttk.Frame(main_frame)
        input_frame.pack(fill=tk.X, pady=(0, 10))
        
        ttk.Label(input_frame, text="文件夹路径:").pack(side=tk.LEFT, padx=(0, 5))
        self.path_var = tk.StringVar()
        self.path_entry = ttk.Entry(input_frame, textvariable=self.path_var, width=80)
        self.path_entry.pack(side=tk.LEFT, padx=(0, 5), fill=tk.X, expand=True)
        
        browse_btn = ttk.Button(input_frame, text="浏览", command=self.browse_folder)
        browse_btn.pack(side=tk.LEFT, padx=(0, 5))
        
        show_btn = ttk.Button(input_frame, text="显示3D树状图", command=self.show_3d_tree)
        show_btn.pack(side=tk.LEFT)
        
        # 缩放控制
        zoom_frame = ttk.Frame(main_frame)
        zoom_frame.pack(fill=tk.X, pady=(0, 5))
        
        ttk.Label(zoom_frame, text="缩放:").pack(side=tk.LEFT, padx=(0, 5))
        self.zoom_var = tk.DoubleVar(value=1.0)
        self.zoom_scale = ttk.Scale(zoom_frame, from_=0.1, to=2.0, variable=self.zoom_var, command=self.on_zoom)
        self.zoom_scale.pack(side=tk.LEFT, fill=tk.X, expand=True, padx=(0, 5))
        
        self.zoom_label = ttk.Label(zoom_frame, text="100%")
        self.zoom_label.pack(side=tk.LEFT)
        
        # 3D显示区域
        self.canvas = Canvas(main_frame, bg="white")
        self.canvas.pack(fill=tk.BOTH, expand=True)
        
        # 缩放和平移支持
        self.scale = 1.0
        self.offset_x = 0
        self.offset_y = 0
        self.last_x = 0
        self.last_y = 0
        self.dragging = False
        
        # 绑定鼠标事件
        self.canvas.bind("<ButtonPress-1>", self.on_press)
        self.canvas.bind("<B1-Motion>", self.on_drag)
        self.canvas.bind("<ButtonRelease-1>", self.on_release)
        self.canvas.bind("<MouseWheel>", self.on_mouse_wheel)
        
        # 文件夹数据
        self.folder_data = {}
        self.max_depth = 0
        self.tree_data = []  # 存储所有节点数据，用于缩放重绘
        
        # 历史记录
        self.history = []
        self.max_history = 3
        self.history_file = "history.txt"
        
        # 加载历史记录
        self.load_history()
        
        # 历史记录显示区域
        history_frame = ttk.Frame(main_frame)
        history_frame.pack(fill=tk.X, pady=(10, 0))
        
        ttk.Label(history_frame, text="历史记录:").pack(side=tk.LEFT, padx=(0, 10))
        self.history_buttons = []
        for i in range(self.max_history):
            btn = ttk.Button(history_frame, text="", width=30, command=lambda idx=i: self.select_history(idx))
            btn.pack(side=tk.LEFT, padx=(0, 5))
            self.history_buttons.append(btn)
        
        # 更新历史记录显示
        self.update_history_display()
        
    def browse_folder(self):
        folder_path = filedialog.askdirectory()
        if folder_path:
            self.path_var.set(folder_path)
            
    def scan_folder(self, folder_path, depth=0):
        if depth > self.max_depth:
            self.max_depth = depth
            
        try:
            items = os.listdir(folder_path)
            folders = []
            files = []
            
            for item in items:
                item_path = os.path.join(folder_path, item)
                if os.path.isdir(item_path):
                    folders.append(item)
                    self.scan_folder(item_path, depth + 1)
                else:
                    files.append(item)
                    
            self.folder_data[folder_path] = {"folders": folders, "files": files, "depth": depth}
            
        except Exception as e:
            print(f"扫描文件夹出错: {e}")
            
    def draw_3d_tree(self):
        self.canvas.delete("all")
        self.tree_data = []
        
        if not self.folder_data:
            return
            
        # 收集树状数据
        self.collect_tree_data(self.path_var.get(), 0, self.canvas.winfo_width()//2, 100)
        
        # 绘制树状图
        self.draw_tree_from_data()
        
    def collect_tree_data(self, folder_path, depth=0, root_x=0, root_y=0):
        if folder_path not in self.folder_data:
            return
            
        data = self.folder_data[folder_path]
        folders = data["folders"]
        files = data["files"]
        
        # 计算当前节点位置
        node = {
            "type": "folder",
            "path": folder_path,
            "name": os.path.basename(folder_path),
            "x": root_x,
            "y": root_y,
            "depth": depth,
            "children": []
        }
        
        self.tree_data.append(node)
        
        # 收集所有层级的节点和父节点关系
        all_nodes = []
        self.collect_all_nodes(folder_path, depth, all_nodes)
        
        # 按层级分组
        level_nodes = {}
        for node_data in all_nodes:
            level = node_data["depth"]
            if level not in level_nodes:
                level_nodes[level] = []
            level_nodes[level].append(node_data)
        
        # 重新计算所有节点位置，使用扇区布局
        self.calculate_sector_positions(root_x, root_y, level_nodes)
        
    def collect_all_nodes(self, folder_path, depth, all_nodes):
        if folder_path not in self.folder_data:
            return
            
        data = self.folder_data[folder_path]
        folders = data["folders"]
        files = data["files"]
        
        # 添加当前文件夹的子节点
        for folder in folders:
            folder_path_full = os.path.join(folder_path, folder)
            node_data = {
                "type": "folder",
                "name": folder,
                "path": folder_path_full,
                "depth": depth + 1,
                "parent": folder_path
            }
            all_nodes.append(node_data)
            self.collect_all_nodes(folder_path_full, depth + 1, all_nodes)
            
        for file in files:
            node_data = {
                "type": "file",
                "name": file,
                "path": os.path.join(folder_path, file),
                "depth": depth + 1,
                "parent": folder_path
            }
            all_nodes.append(node_data)
            
    def calculate_sector_positions(self, root_x, root_y, level_nodes):
        # 清空原有数据，重新绘制
        self.tree_data = [{
            "type": "folder",
            "name": os.path.basename(self.path_var.get()),
            "path": self.path_var.get(),
            "x": root_x,
            "y": root_y,
            "depth": 0
        }]
        
        # 先计算所有节点的位置
        all_node_data = []
        
        # 为每个层级计算位置
        for level in sorted(level_nodes.keys()):
            nodes = level_nodes[level]
            total_nodes = len(nodes)
            if total_nodes == 0:
                continue
                
            radius = level * 100
            
            # 按父节点分组
            parent_groups = {}
            for node_data in nodes:
                parent = node_data["parent"]
                if parent not in parent_groups:
                    parent_groups[parent] = []
                parent_groups[parent].append(node_data)
            
            # 为每个父节点的子节点计算位置
            for parent_path, group_nodes in parent_groups.items():
                parent_node = next((n for n in self.tree_data if n["path"] == parent_path), None)
                if not parent_node:
                    continue
                    
                # 计算当前层级的扇区范围
                # 父节点所在环的文件夹个数
                parent_level = parent_node["depth"]
                parent_level_nodes = []
                if parent_level in level_nodes:
                    parent_level_nodes = [n for n in level_nodes[parent_level] if n["type"] == "folder"]
                else:
                    parent_level_nodes = [parent_node] if parent_node["type"] == "folder" else []
                    
                total_parent_folders = len(parent_level_nodes)
                if total_parent_folders == 0:
                    total_parent_folders = 1
                    
                # 每个文件夹分配360度/total_parent_folders的扇区
                sector_angle = 2 * math.pi / total_parent_folders
                
                # 找到父节点在同级中的位置
                parent_index = 0
                if parent_level in level_nodes and parent_node["type"] == "folder":
                    for i, n in enumerate(parent_level_nodes):
                        if n["path"] == parent_path:
                            parent_index = i
                            break
                
                # 计算父节点的扇区
                start_angle = parent_index * sector_angle
                end_angle = (parent_index + 1) * sector_angle
                sector_width = end_angle - start_angle
                
                # 在扇区内均匀分布子节点
                angle_step = sector_width / len(group_nodes)
                
                for i, node_data in enumerate(group_nodes):
                    angle = start_angle + (i + 0.5) * angle_step
                    # 文件夹在大圆上，文件在小圆上
                    if node_data["type"] == "folder":
                        current_radius = radius
                    else:
                        current_radius = radius - 40
                    
                    x = root_x + current_radius * math.cos(angle)
                    y = root_y + current_radius * math.sin(angle)
                    
                    # 保存扇区信息
                    node_data["sector_start"] = start_angle + i * angle_step
                    node_data["sector_end"] = start_angle + (i + 1) * angle_step
                    node_data["x"] = x
                    node_data["y"] = y
                    node_data["angle"] = angle
                    
                    all_node_data.append(node_data)
                    
                    self.tree_data.append({
                        "type": node_data["type"],
                        "name": node_data["name"],
                        "path": node_data["path"],
                        "x": x,
                        "y": y,
                        "depth": level,
                        "parent_path": node_data["parent"],
                        "sector_start": node_data["sector_start"],
                        "sector_end": node_data["sector_end"],
                        "angle": angle
                    })
                    
                    # 保存扇区信息
                    node_data["sector_start"] = start_angle + i * angle_step
                    node_data["sector_end"] = start_angle + (i + 1) * angle_step
        
        # 调整父节点位置，使其与子节点中间位置对齐
        # 从最深层级开始向上调整
        max_level = max(level_nodes.keys()) if level_nodes else 0
        
        for level in range(max_level, 0, -1):
            if level not in level_nodes:
                continue
                
            # 找到当前层级的文件夹
            current_folders = [n for n in level_nodes[level] if n["type"] == "folder"]
            
            for folder in current_folders:
                # 找到该文件夹的所有子节点
                child_nodes = [n for n in all_node_data if n["parent"] == folder["path"]]
                if not child_nodes:
                    continue
                    
                # 计算子节点的平均角度
                total_angle = sum(n["angle"] for n in child_nodes)
                avg_angle = total_angle / len(child_nodes)
                
                # 找到父节点
                parent_node = next((n for n in self.tree_data if n["path"] == folder["path"]), None)
                if parent_node:
                    # 调整父节点位置
                    radius = level * 100
                    x = root_x + radius * math.cos(avg_angle)
                    y = root_y + radius * math.sin(avg_angle)
                    parent_node["x"] = x
                    parent_node["y"] = y
                    parent_node["angle"] = avg_angle
                
    def draw_node(self, x, y, folder_path, depth, width, level_spacing):
        # 改为使用收集的数据进行绘制
        pass
        
    def draw_tree_from_data(self):
        self.canvas.delete("all")
        
        for node in self.tree_data:
            # 应用缩放和平移
            x = node["x"] * self.scale + self.offset_x
            y = node["y"] * self.scale + self.offset_y
            
            if node["type"] == "folder":
                node_radius = min(15, 15 * self.scale)  # 缩小文件夹节点
                self.canvas.create_oval(x - node_radius, y - node_radius, x + node_radius, y + node_radius, fill="#4CAF50", outline="#2E7D32")
                font_size = min(8, int(8 * self.scale))  # 缩小字体
                self.canvas.create_text(x, y, text=node["name"][:8] + "..." if len(node["name"]) > 8 else node["name"], fill="white", font=("Arial", font_size, "bold"))
            else:
                file_width = min(20, 20 * self.scale)  # 缩小文件节点
                file_height = min(15, 15 * self.scale)
                self.canvas.create_rectangle(x - file_width//2, y - file_height//2, x + file_width//2, y + file_height//2, fill="#2196F3", outline="#1565C0")
                font_size = min(6, int(6 * self.scale))  # 缩小字体
                self.canvas.create_text(x, y, text=node["name"][:6] + "..." if len(node["name"]) > 6 else node["name"], fill="white", font=("Arial", font_size))
                
        # 绘制连接线
        for node in self.tree_data:
            if "parent_path" in node:
                # 找到父节点
                parent_node = next((n for n in self.tree_data if n["path"] == node["parent_path"]), None)
                if parent_node:
                    x1 = parent_node["x"] * self.scale + self.offset_x
                    y1 = parent_node["y"] * self.scale + self.offset_y
                    x2 = node["x"] * self.scale + self.offset_x
                    y2 = node["y"] * self.scale + self.offset_y
                    self.canvas.create_line(x1, y1, x2, y2, fill="#666666", width=1)
        
        # 绘制从圆心到文件夹的红色射线，向外射出
        root_node = self.tree_data[0]
        # 使用未缩放的坐标进行计算
        root_x_unscaled = root_node["x"]
        root_y_unscaled = root_node["y"]
        
        for node in self.tree_data:
            if node["type"] == "folder" and node["depth"] > 0:  # 排除根节点
                # 使用未缩放的坐标进行计算
                folder_x_unscaled = node["x"]
                folder_y_unscaled = node["y"]
                
                # 计算射线方向向量（未缩放）
                dx_unscaled = folder_x_unscaled - root_x_unscaled
                dy_unscaled = folder_y_unscaled - root_y_unscaled
                
                # 计算向量长度（未缩放）
                distance_unscaled = math.sqrt(dx_unscaled * dx_unscaled + dy_unscaled * dy_unscaled)
                if distance_unscaled > 0:
                    # 归一化方向向量（未缩放）
                    unit_dx = dx_unscaled / distance_unscaled
                    unit_dy = dy_unscaled / distance_unscaled
                    
                    # 计算射线延伸的终点（未缩放）
                    extension_length_unscaled = 100  # 延伸长度，100像素（未缩放）
                    end_x_unscaled = folder_x_unscaled + unit_dx * extension_length_unscaled
                    end_y_unscaled = folder_y_unscaled + unit_dy * extension_length_unscaled
                    
                    # 应用缩放和偏移，准备绘制
                    root_x = root_x_unscaled * self.scale + self.offset_x
                    root_y = root_y_unscaled * self.scale + self.offset_y
                    folder_x = folder_x_unscaled * self.scale + self.offset_x
                    folder_y = folder_y_unscaled * self.scale + self.offset_y
                    end_x = end_x_unscaled * self.scale + self.offset_x
                    end_y = end_y_unscaled * self.scale + self.offset_y
                    
                    # 绘制红色射线，从圆心到延伸的终点
                    self.canvas.create_line(root_x, root_y, end_x, end_y, fill="#FF0000", width=1)
                    
                    # 计算红色射线与文件所在圆圈的交点，并绘制绿色点
                    # 文件所在圆圈的半径（基于文件的深度，未缩放）
                    # 文件的深度是文件夹深度 + 1
                    folder_depth = node["depth"]
                    file_depth = folder_depth + 1
                    file_circle_radius_unscaled = file_depth * 100 - 40
                    
                    # 计算射线与文件所在圆圈的交点（使用未缩放的坐标）
                    # 使用勾股定理计算交点
                    a = dx_unscaled * dx_unscaled + dy_unscaled * dy_unscaled
                    b = 2 * (root_x_unscaled * dx_unscaled + root_y_unscaled * dy_unscaled)
                    c = root_x_unscaled * root_x_unscaled + root_y_unscaled * root_y_unscaled - file_circle_radius_unscaled * file_circle_radius_unscaled
                    
                    # 计算判别式
                    discriminant = b * b - 4 * a * c
                    if discriminant >= 0:
                        # 计算两个交点参数
                        t1 = (-b + math.sqrt(discriminant)) / (2 * a)
                        t2 = (-b - math.sqrt(discriminant)) / (2 * a)
                        
                        # 找到在射线方向上的交点（t>0）
                        for t in [t1, t2]:
                            if t > 0:
                                # 计算交点（未缩放）
                                intersection_x_unscaled = root_x_unscaled + t * dx_unscaled
                                intersection_y_unscaled = root_y_unscaled + t * dy_unscaled
                                # 应用缩放和偏移
                                intersection_x = intersection_x_unscaled * self.scale + self.offset_x
                                intersection_y = intersection_y_unscaled * self.scale + self.offset_y
                                # 绘制绿色点
                                self.canvas.create_oval(intersection_x - 3, intersection_y - 3, intersection_x + 3, intersection_y + 3, fill="#00FF00", outline="#00FF00", width=1)
        
        # 绘制同心圆
        max_level = max(node["depth"] for node in self.tree_data)
        # Get root node position after scale and offset
        root_node = self.tree_data[0]
        root_x = root_node["x"] * self.scale + self.offset_x
        root_y = root_node["y"] * self.scale + self.offset_y
        for level in range(1, max_level + 1):
            radius = level * 100 * self.scale
            self.canvas.create_oval(root_x - radius, root_y - radius, root_x + radius, root_y + radius, outline="#333333", width=1)
            # 在内侧40px处添加小同心圆
            inner_radius = radius - 40 * self.scale
            if inner_radius > 0:
                self.canvas.create_oval(root_x - inner_radius, root_y - inner_radius, root_x + inner_radius, root_y + inner_radius, outline="#666666", width=1)
                
    def draw_3d_line(self, x1, y1, x2, y2):
        # 绘制连接线，使用更细的线条
        self.canvas.create_line(x1, y1, x2, y2, fill="#757575", width=1)
        self.canvas.create_line(x1 + 1, y1 + 1, x2 + 1, y2 + 1, fill="#BDBDBD", width=0.5)
        
    def on_zoom(self, value):
        self.scale = float(value)
        self.zoom_label.config(text=f"{int(self.scale * 100)}%")
        if self.tree_data:
            self.draw_tree_from_data()
            
    def on_press(self, event):
        self.last_x = event.x
        self.last_y = event.y
        self.dragging = True
        
    def on_drag(self, event):
        if self.dragging:
            dx = event.x - self.last_x
            dy = event.y - self.last_y
            self.offset_x += dx
            self.offset_y += dy
            self.last_x = event.x
            self.last_y = event.y
            if self.tree_data:
                self.draw_tree_from_data()
                
    def on_release(self, event):
        self.dragging = False
        
    def on_mouse_wheel(self, event):
        # 鼠标滚轮缩放
        delta = event.delta
        if delta > 0:
            self.scale = min(2.0, self.scale * 1.1)
        else:
            self.scale = max(0.1, self.scale * 0.9)
            
        self.zoom_var.set(self.scale)
        self.zoom_label.config(text=f"{int(self.scale * 100)}%")
        if self.tree_data:
            self.draw_tree_from_data()
        
    def show_3d_tree(self):
        folder_path = self.path_var.get()
        if not folder_path or not os.path.isdir(folder_path):
            messagebox.showerror("错误", "请选择有效的文件夹路径")
            return
        
        # 添加到历史记录
        self.add_to_history(folder_path)
            
        self.folder_data = {}
        self.max_depth = 0
        self.scan_folder(folder_path)
        self.draw_3d_tree()
    
    def add_to_history(self, folder_path):
        # 如果路径已在历史记录中，先移除
        if folder_path in self.history:
            self.history.remove(folder_path)
        
        # 添加到历史记录开头
        self.history.insert(0, folder_path)
        
        # 限制历史记录数量
        if len(self.history) > self.max_history:
            self.history = self.history[:self.max_history]
        
        # 更新历史记录显示
        self.update_history_display()
        
        # 保存历史记录到文件
        self.save_history()
    
    def update_history_display(self):
        # 更新历史记录按钮
        for i in range(self.max_history):
            if i < len(self.history):
                # 显示路径的最后部分
                path_parts = self.history[i].split(os.sep)
                display_text = os.sep.join(path_parts[-3:])
                if len(display_text) > 28:
                    display_text = "..." + display_text[-25:]
                self.history_buttons[i].config(text=display_text, state=tk.NORMAL)
            else:
                self.history_buttons[i].config(text="", state=tk.DISABLED)
    
    def select_history(self, index):
        if index < len(self.history):
            folder_path = self.history[index]
            self.path_var.set(folder_path)
            # 重新显示树状图
            self.show_3d_tree()
    
    def load_history(self):
        """加载历史记录"""
        try:
            if os.path.exists(self.history_file):
                with open(self.history_file, 'r', encoding='utf-8') as f:
                    lines = f.readlines()
                    self.history = [line.strip() for line in lines if line.strip()]
                    # 限制历史记录数量
                    if len(self.history) > self.max_history:
                        self.history = self.history[:self.max_history]
        except Exception as e:
            print(f"加载历史记录出错: {e}")
    
    def save_history(self):
        """保存历史记录"""
        try:
            with open(self.history_file, 'w', encoding='utf-8') as f:
                for path in self.history:
                    f.write(path + '\n')
        except Exception as e:
            print(f"保存历史记录出错: {e}")

if __name__ == "__main__":
    root = tk.Tk()
    app = FolderTree3D(root)
    root.mainloop()