빗자루네

# vi: set sw=4 sts=4 expandtab:
import wx
import os

class wxPathTreeCtrl(wx.TreeCtrl):
    def __init__(self, *args, **kwargs):
        wx.TreeCtrl.__init__(self, *args, **kwargs) 

    def InitRoot(self, rootPath):
        self.branchd = {}
        self.DeleteAllItems()
        self.root = self.AddRoot(rootPath)
        self.SetPyData(self.root, ("ROOT", rootPath))
        return self.root

    def AppendPath(self, path):
        branch = path

        branches = []
        while branch:
            branch, leaf = os.path.split(branch)
            if branch in self.branchd:
                break
            elif branch:
                branches.append(branch)

        branches.reverse()

        last = self.branchd.get(branch, self.root)
        for branch in branches:
            last = self.AppendItem(last, os.path.split(branch)[1])
            self.SetPyData(last, ("DIR", branch))
            self.branchd[branch] = last

        
        item = self.AppendItem(last, os.path.split(path)[1])
        self.SetPyData(item, ("FILE", path))
        return item

    def ExpandAllDirs(self):
        self.ExpandDirs(self.root)

    def ExpandDirs(self, node):
        if self.GetPyData(node)[0] == "FILE":
            pass
        else:
            children = list(self.GenChildren(node))
            fileCount = len([child for child in children if self.GetPyData(child)[0] == "FILE"])
            if 0 == fileCount:
                self.Expand(node)

                for child in children:
                    self.ExpandDirs(child)

    def GetBranchd(self):
        return self.branchd

    def GenChildren(self, node):
        child = self.GetFirstChild(node)[0]
        while child:
            yield child
            child = self.GetNextSibling(child)

if __name__ == "__main__":

    class TestFrame(wx.Frame):
        def __init__(self, parent, title, size=(800, 600)):
            wx.Frame.__init__(self, parent, -1, title, pos=(0, 0), size=size)
            self.CentreOnScreen(wx.BOTH)

            self.pathTreeCtrl = wxPathTreeCtrl(self)
            self.pathTreeCtrl.InitRoot("root")
            self.pathTreeCtrl.AppendPath("bin/main.exe")
            self.pathTreeCtrl.AppendPath("data/char/pc/warrior/warriror.png")
            self.pathTreeCtrl.AppendPath("data/char/pc/warrior/warriror2.png")
            self.pathTreeCtrl.ExpandAllDirs()

    class TestApp(wx.App):
        def OnInit(self):
            toolFrame = TestFrame(None, "TEST")
            toolFrame.Show()
            self.SetTopWindow(toolFrame)
            
            return True

    TestApp(redirect=False).MainLoop()