IamHere/folder_tree_3d.py

import tkinter as tk
from tkinter import filedialog, ttk
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 = []  # 存储所有节点数据，用于缩放重绘
        
    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 - 20
                    
                    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 = root_node["x"] * self.scale + self.offset_x
        root_y = root_node["y"] * self.scale + self.offset_y
        
        for node in self.tree_data:
            if node["type"] == "file":
                # 计算根节点与文件节点连线和文件所在圆的交点
                file_x = node["x"] * self.scale + self.offset_x
                file_y = node["y"] * self.scale + self.offset_y
                
                # 计算文件所在圆的半径
                file_depth = node["depth"]
                file_radius = (file_depth * 100 - 20) * self.scale
                
                # 计算交点
                dx = file_x - root_x
                dy = file_y - root_y
                distance = math.sqrt(dx * dx + dy * dy)
                if distance == 0:
                    continue
                
                # 计算交点坐标
                intersection_x = root_x + (dx / distance) * file_radius
                intersection_y = root_y + (dy / distance) * file_radius
                
                # 标记交点
                self.canvas.create_oval(intersection_x - 3, intersection_y - 3, intersection_x + 3, intersection_y + 3, fill="#FF0000", outline="#FF0000", 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)
            # 在内侧20px处添加小同心圆
            inner_radius = radius - 20 * 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):
            tk.messagebox.showerror("错误", "请选择有效的文件夹路径")
            return
            
        self.folder_data = {}
        self.max_depth = 0
        self.scan_folder(folder_path)
        self.draw_3d_tree()

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